Sampling

src/pyjuice/layer/input_layer.py

@@ -338,6 +338,7 @@ def forward(self, data: torch.Tensor, node_mars: torch.Tensor, params: Optional[
                 if not cond_fn(self, kwargs):
                     continue
 
+                kwargs["batch_size"] = batch_size
                 target_kwargs, grid = prep_kwargs_fn(self, kwargs)
 
                 if grid is None:
@@ -487,6 +488,7 @@ def backward(self, data: torch.Tensor, node_flows: torch.Tensor,
                 if not cond_fn(self, kwargs):
                     continue
 
+                kwargs["batch_size"] = batch_size
                 target_kwargs, grid = prep_kwargs_fn(self, kwargs)
 
                 if grid is None:
@@ -562,7 +564,7 @@ def backward(self, data: torch.Tensor, node_flows: torch.Tensor,
             raise NotImplementedError("CPU backward fn for input nodes is not implemented.")
 
     def sample(self, samples: torch.Tensor, node_flows: torch.Tensor, missing_mask: Optional[torch.Tensor] = None, 
-               params: Optional[torch.Tensor] = None, seed: Optional[int] = None):
+               params: Optional[torch.Tensor] = None, seed: Optional[int] = None, **kwargs):
         """
         samples:       [num_vars, B]
         missing_mask:  [num_vars, B] or [num_vars] or None
@@ -576,43 +578,80 @@ def sample(self, samples: torch.Tensor, node_flows: torch.Tensor, missing_mask:
 
         assert params.dim() == 1
 
-        if "cuda" in self.device.type:
+        if len(self.get_dist().sampling_fns) == 0:
+            if "cuda" in self.device.type:
 
-            sid, eid = self._output_ind_range
-            tot_num_nodes = node_flows.size(0)
-            batch_size = node_flows.size(1)
-            node_offset = self._output_ind_range[0]
+                sid, eid = self._output_ind_range
+                tot_num_nodes = node_flows.size(0)
+                batch_size = node_flows.size(1)
+                node_offset = self._output_ind_range[0]
 
-            # Get all node ids with non-zero flow
-            nflow_xids, nflow_yids = torch.where(node_flows[sid:eid,:] > 1e-8)
-            num_activ_nodes = nflow_xids.size(0)
+                # Get all node ids with non-zero flow
+                nflow_xids, nflow_yids = torch.where(node_flows[sid:eid,:] > 1e-8)
+                num_activ_nodes = nflow_xids.size(0)
 
-            if not self.provided("_sample_kernel"):
-                self._sample_kernel = self._compile_triton_kernel(self._sample_kernel_template, sample_fn = self.sample_fn)
+                if not self.provided("_sample_kernel"):
+                    self._sample_kernel = self._compile_triton_kernel(self._sample_kernel_template, sample_fn = self.sample_fn)
 
-            BLOCK_SIZE = 1024
+                BLOCK_SIZE = 1024
 
-            grid = (triton.cdiv(num_activ_nodes, BLOCK_SIZE),)
+                grid = (triton.cdiv(num_activ_nodes, BLOCK_SIZE),)
 
-            self._sample_kernel[grid](
-                samples_ptr = samples, 
-                params_ptr = params,
-                nflow_xids_ptr = nflow_xids, 
-                nflow_yids_ptr = nflow_yids, 
-                vids_ptr = self.vids, 
-                s_pids_ptr = self.s_pids, 
-                metadata_ptr = self.metadata,
-                s_mids_ptr = self.s_mids,
-                num_activ_nodes = num_activ_nodes, 
-                num_vars_per_node = self.num_vars_per_node, 
-                nv_block_size = triton.next_power_of_2(self.num_vars_per_node),
-                batch_size = batch_size, 
-                BLOCK_SIZE = BLOCK_SIZE,
-                seed = seed if seed is not None else random.randint(0, int(1e8))
-            )
+                self._sample_kernel[grid](
+                    samples_ptr = samples, 
+                    params_ptr = params,
+                    nflow_xids_ptr = nflow_xids, 
+                    nflow_yids_ptr = nflow_yids, 
+                    vids_ptr = self.vids, 
+                    s_pids_ptr = self.s_pids, 
+                    metadata_ptr = self.metadata,
+                    s_mids_ptr = self.s_mids,
+                    num_activ_nodes = num_activ_nodes, 
+                    num_vars_per_node = self.num_vars_per_node, 
+                    nv_block_size = triton.next_power_of_2(self.num_vars_per_node),
+                    batch_size = batch_size, 
+                    BLOCK_SIZE = BLOCK_SIZE,
+                    seed = seed if seed is not None else random.randint(0, int(1e8))
+                )
+
+            else:
+                raise NotImplementedError("CPU sample fn for input nodes is not implemented.")
 
         else:
-            raise NotImplementedError("CPU sample fn for input nodes is not implemented.")
+            for (kernel, cond_fn, prep_kwargs_fn) in self.get_dist().sampling_fns:
+                if not cond_fn(self, kwargs):
+                    continue
+
+                batch_size = node_flows.size(1)
+
+                # Get all node ids with non-zero flow
+                sid, eid = self._output_ind_range
+                nflow_xids, nflow_yids = torch.where(node_flows[sid:eid,:] > 1e-8)
+                num_activ_nodes = nflow_xids.size(0)
+
+                kwargs["num_activ_nodes"] = num_activ_nodes
+                kwargs["batch_size"] = batch_size
+                target_kwargs, grid = prep_kwargs_fn(self, kwargs)
+
+                if grid is None:
+                    grid = (triton.cdiv(num_activ_nodes, target_kwargs["BLOCK_SIZE"]),)
+
+                kernel[grid](
+                    samples_ptr = samples, 
+                    params_ptr = params,
+                    nflow_xids_ptr = nflow_xids, 
+                    nflow_yids_ptr = nflow_yids, 
+                    vids_ptr = self.vids, 
+                    s_pids_ptr = self.s_pids, 
+                    metadata_ptr = self.metadata,
+                    s_mids_ptr = self.s_mids,
+                    num_activ_nodes = num_activ_nodes, 
+                    num_vars_per_node = self.num_vars_per_node, 
+                    nv_block_size = triton.next_power_of_2(self.num_vars_per_node),
+                    batch_size = batch_size,
+                    seed = seed if seed is not None else random.randint(0, int(1e8)),
+                    **target_kwargs
+                )
 
     def mini_batch_em(self, step_size: float, pseudocount: float = 0.0, keep_zero_params: bool = False):
         if not self._used_external_params:

src/pyjuice/nodes/distributions/categorical.py

@@ -141,7 +141,7 @@ def sample_fn(samples_ptr, local_offsets, batch_offsets, vids, s_pids, params_pt
 
         max_num_cats = tl.max(num_cats, axis = 0)
 
-        rnd_val = tl.rand(seed, tl.arange(0, BLOCK_SIZE))
+        rnd_val = tl.rand(seed, tl.arange(0, BLOCK_SIZE) + block_start)
         sampled_id = tl.zeros([BLOCK_SIZE], dtype = tl.int64) - 1
 
         # Sample by computing cumulative probability

src/pyjuice/nodes/distributions/distributions.py

@@ -9,6 +9,7 @@ class Distribution():
     def __init__(self):
         self.post_fw_fns = []
         self.post_bp_fns = []
+        self.sampling_fns = []
 
     def get_signature(self):
         """

src/pyjuice/nodes/distributions/external_categorical.py

@@ -63,6 +63,8 @@ def _prep_args_apply_ll_kernel(layer, kwargs):
     else:
         raise ValueError("Unexpected `extern_product_categorical_mode`. Should be 'normalized_ll', 'unnormalized_ll', or 'normalizing_constant'.")
 
+    assert kwargs["batch_size"] == external_categorical_logps.size(0), "Batch size doesn't match in `external_categorical_logps`."
+
     target_kwargs["compute_unnorm_logp"] = compute_unnorm_logp
     target_kwargs["compute_logz"] = compute_logz
     target_kwargs["ext_softevi_indexing"] = ext_softevi_indexing
@@ -116,6 +118,8 @@ def _prep_args_apply_ll_bp_kernel(layer, kwargs):
     external_categorical_logps = kwargs["external_categorical_logps"]
     assert external_categorical_logps.dim() == 3
 
+    assert kwargs["batch_size"] == external_categorical_logps.size(0), "Batch size doesn't match in `external_categorical_logps`."
+
     target_kwargs["external_categorical_logps_ptr"] = external_categorical_logps
 
     target_kwargs["var_idmapping_ptr"] = layer.var_idmapping
@@ -176,6 +180,8 @@ def _prep_args_apply_ll_bp_extern_grad_kernel(layer, kwargs):
     external_categorical_logps = kwargs["external_categorical_logps"]
     assert external_categorical_logps.dim() == 3
 
+    assert kwargs["batch_size"] == external_categorical_logps.size(0), "Batch size doesn't match in `external_categorical_logps`."
+
     assert "external_categorical_logps_grad" in kwargs
     external_categorical_logps_grad = kwargs["external_categorical_logps_grad"]
     assert external_categorical_logps_grad.shape == external_categorical_logps.shape
@@ -214,6 +220,36 @@ def _prep_args_apply_ll_bp_extern_grad_kernel(layer, kwargs):
     return target_kwargs, grid
 
 
+def _condition_sample_kernel(layer, kwargs):
+    return "external_categorical_logps" in kwargs and kwargs["external_categorical_logps"].dim() == 3
+
+
+def _prep_args_sample_kernel(layer, kwargs):
+    target_kwargs = dict()
+
+    external_categorical_logps = kwargs["external_categorical_logps"]
+
+    assert kwargs["batch_size"] == external_categorical_logps.size(0), "Batch size doesn't match."
+
+    target_kwargs["external_categorical_logps_ptr"] = external_categorical_logps
+
+    target_kwargs["var_idmapping_ptr"] = layer.var_idmapping
+
+    target_kwargs["ext_num_vars"] = external_categorical_logps.size(1)
+
+    target_kwargs["max_num_cats"] = external_categorical_logps.size(2)
+
+    num_activ_nodes = kwargs["num_activ_nodes"]
+
+    target_kwargs["TILE_SIZE_K"] = min(64, triton.next_power_of_2(target_kwargs["max_num_cats"]))
+    target_kwargs["K_NUM_TILES"] = triton.cdiv(target_kwargs["max_num_cats"], target_kwargs["TILE_SIZE_K"])
+    target_kwargs["BLOCK_S"] = min(64, 1024 // target_kwargs["TILE_SIZE_K"], triton.next_power_of_2(num_activ_nodes))
+
+    grid = (triton.cdiv(num_activ_nodes, target_kwargs["BLOCK_S"]),)
+
+    return target_kwargs, grid
+
+
 class ExternProductCategorical(Distribution):
     """
     A class representing a product distribution of two Categorical distributions: Pr_{A} * Pr_{B}.
@@ -237,6 +273,10 @@ def __init__(self, num_cats: int):
             (self.ll_bp_extern_grad_kernel, _condition_apply_ll_bp_extern_grad_kernel, _prep_args_apply_ll_bp_extern_grad_kernel)
         ]
 
+        self.sampling_fns = [
+            (self.sample_kernel, _condition_sample_kernel, _prep_args_sample_kernel)
+        ]
+
     def get_signature(self):
         """
         Get the signature of the current distribution.
@@ -591,6 +631,88 @@ def ll_bp_extern_grad_kernel(params_ptr, param_flows_ptr, node_flows_ptr, node_m
 
                 tl.atomic_add(expars_grad_ptr + i * TILE_SIZE_K, parflows, mask = cat_mask_ex)
 
+    @staticmethod
+    @triton_jit
+    def sample_kernel(samples_ptr, params_ptr, nflow_xids_ptr, nflow_yids_ptr, vids_ptr, s_pids_ptr, metadata_ptr, s_mids_ptr,
+                      num_activ_nodes, num_vars_per_node: tl.constexpr, nv_block_size: tl.constexpr, batch_size: tl.constexpr, seed, 
+                      external_categorical_logps_ptr, var_idmapping_ptr, ext_num_vars: tl.constexpr, max_num_cats: tl.constexpr,
+                      TILE_SIZE_K: tl.constexpr, K_NUM_TILES: tl.constexpr, BLOCK_S: tl.constexpr):
+        pid = tl.program_id(axis = 0)
+        block_start = pid * BLOCK_S
+
+        offsets = block_start + tl.arange(0, BLOCK_S) # [BLOCK_S]
+        mask = offsets < num_activ_nodes
+
+        # Raw batch and (local) node id
+        local_offsets = tl.load(nflow_xids_ptr + offsets, mask = mask, other = 0)
+        batch_offsets = tl.load(nflow_yids_ptr + offsets, mask = mask, other = 0)
+
+        # Load variable ids from `vids_ptr`
+        vids = tl.load(vids_ptr + local_offsets, mask = mask, other = 0)
+        lvids = tl.load(var_idmapping_ptr + vids, mask = mask, other = 0)
+
+        # Get `num_cats` from `metadata`
+        s_mids = tl.load(s_mids_ptr + local_offsets, mask = mask, other = 0)
+        num_cats = tl.load(metadata_ptr + s_mids, mask = mask, other = 0).to(tl.int64) # [BLOCK_SIZE]
+
+        # Get start parameter indices
+        s_pids = tl.load(s_pids_ptr + local_offsets, mask = mask, other = 0)
+
+        # Ptrs pointing to internal parameters
+        inpars_ptr = params_ptr + s_pids[:,None] + tl.arange(0, TILE_SIZE_K)[None,:] # [BLOCK_S, TILE_SIZE_K]
+
+        # Ptrs pointing to external parameters
+        expars_ptr = external_categorical_logps_ptr + \
+            batch_offsets[:,None] * (ext_num_vars * max_num_cats) + \
+            lvids[:,None] * max_num_cats + \
+            tl.arange(0, TILE_SIZE_K)[None,:]  # [BLOCK_S, TILE_SIZE_K]
+
+        # Compute logZ
+        logZ = tl.zeros([BLOCK_S], dtype = tl.float32) - float("inf")
+        for i in range(K_NUM_TILES):
+            cat_mask = mask[:,None] & (i * TILE_SIZE_K + tl.arange(0, TILE_SIZE_K)[None,:] < num_cats[:,None])
+
+            inpar = tl.load(inpars_ptr + i * TILE_SIZE_K, mask = cat_mask, other = 0.0)
+            expar = tl.load(expars_ptr + i * TILE_SIZE_K, mask = cat_mask, other = 0.0)
+
+            addlpar = inpar.log() + expar
+            addlpar_max = tl.max(addlpar, axis = 1)
+            lpar = (addlpar - addlpar_max[:,None]).exp().sum(axis = 1).log() + addlpar_max
+
+            # Compute log-add-exp(logZ, lpar)
+            maxval = tl.maximum(logZ, lpar)
+            minval = tl.minimum(logZ, lpar)
+            diff = minval - maxval
+
+            logZ = tl.where(logZ == -float("inf"),
+                lpar,
+                maxval + tlmath.log1p(tl.exp(diff))
+            )
+
+        # Generate random number
+        rnd_val = tl.rand(seed, offsets)
+
+        # Draw samples
+        sampled_ids = tl.zeros([BLOCK_S], dtype = tl.int64) - 1
+        for i in range(K_NUM_TILES):
+            cat_mask = mask[:,None] & (i * TILE_SIZE_K + tl.arange(0, TILE_SIZE_K)[None,:] < num_cats[:,None])
+
+            inpar = tl.load(inpars_ptr + i * TILE_SIZE_K, mask = cat_mask, other = 0.0)
+            expar = tl.load(expars_ptr + i * TILE_SIZE_K, mask = cat_mask, other = 0.0)
+
+            probs = tl.exp(tl.log(inpar) + expar - logZ[:,None]) # [BLOCK_S, TILE_SIZE_K]
+            cum_probs = tl.cumsum(probs, axis = 1) # [BLOCK_S, TILE_SIZE_K]
+
+            local_catids = tl.sum((rnd_val[:,None] >= cum_probs).to(tl.int64), axis = 1) # [BLOCK_S]
+
+            is_overflow = (local_catids == TILE_SIZE_K)
+            rnd_val = tl.where(is_overflow, rnd_val - tl.sum(probs, axis = 1), rnd_val)
+            sampled_ids = tl.where(is_overflow | (sampled_ids > -1), sampled_ids, local_catids + i * TILE_SIZE_K)
+
+        # Write back to `samples`
+        sample_offsets = vids * batch_size + batch_offsets
+        tl.store(samples_ptr + sample_offsets, sampled_ids, mask = mask)
+
     @staticmethod
     def bk_flow_mask_fn(local_offsets, ns_offsets, data, flows, node_mars_ptr, params_ptr, param_flows_ptr, s_pids, s_pfids, metadata_ptr, 
                         s_mids_ptr, mask, num_vars_per_node, BLOCK_SIZE, TILE_SIZE_K):