pyRDDLGym-jax 0.2__py3-none-any.whl → 0.4__py3-none-any.whl

pyRDDLGym_jax/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.4'

pyRDDLGym_jax/core/compiler.py

@@ -1,23 +1,11 @@
 from functools import partial
+import traceback
+from typing import Any, Callable, Dict, List, Optional
+
 import jax
 import jax.numpy as jnp
 import jax.random as random
 import jax.scipy as scipy 
-import traceback
-from typing import Any, Callable, Dict, List, Optional
-
-from pyRDDLGym.core.debug.exception import raise_warning
-
-# more robust approach - if user does not have this or broken try to continue
-try:
-    from tensorflow_probability.substrates import jax as tfp
-except Exception:
-    raise_warning('Failed to import tensorflow-probability: '
-                  'compilation of some complex distributions '
-                  '(Binomial, Negative-Binomial, Multinomial) will fail. '
-                  'Please ensure this package is installed correctly.', 'red')
-    traceback.print_exc()
-    tfp = None
 
 from pyRDDLGym.core.compiler.initializer import RDDLValueInitializer
 from pyRDDLGym.core.compiler.levels import RDDLLevelAnalysis
@@ -26,12 +14,23 @@ from pyRDDLGym.core.compiler.tracer import RDDLObjectsTracer
 from pyRDDLGym.core.constraints import RDDLConstraints
 from pyRDDLGym.core.debug.exception import (
     print_stack_trace, 
+    raise_warning,
     RDDLInvalidNumberOfArgumentsError, 
     RDDLNotImplementedError
 )
 from pyRDDLGym.core.debug.logger import Logger
 from pyRDDLGym.core.simulator import RDDLSimulatorPrecompiled
 
+# more robust approach - if user does not have this or broken try to continue
+try:
+    from tensorflow_probability.substrates import jax as tfp
+except Exception:
+    raise_warning('Failed to import tensorflow-probability: '
+                  'compilation of some complex distributions '
+                  '(Binomial, Negative-Binomial, Multinomial) will fail.', 'red')
+    traceback.print_exc()
+    tfp = None
+
 
 # ===========================================================================
 # EXACT RDDL TO JAX COMPILATION RULES
@@ -88,7 +87,7 @@ def _function_aggregation_exact_named(op, name):
 def _function_if_exact_named():
         
     def _jax_wrapped_if_exact(c, a, b, param):
-        return jnp.where(c, a, b)
+        return jnp.where(c > 0.5, a, b)
         
     return _jax_wrapped_if_exact
     
@@ -115,16 +114,27 @@ def _function_bernoulli_exact_named():
 def _function_discrete_exact_named():
         
     def _jax_wrapped_discrete_exact(key, prob, param):
-        logits = jnp.log(prob)
-        sample = random.categorical(key=key, logits=logits, axis=-1)
-        out_of_bounds = jnp.logical_not(jnp.logical_and(
-            jnp.all(prob >= 0),
-            jnp.allclose(jnp.sum(prob, axis=-1), 1.0)))
-        return sample, out_of_bounds
+        return random.categorical(key=key, logits=jnp.log(prob), axis=-1)
         
     return _jax_wrapped_discrete_exact
 
 
+def _function_poisson_exact_named():
+    
+    def _jax_wrapped_poisson_exact(key, rate, param):
+        return random.poisson(key=key, lam=rate, dtype=jnp.int64)
+    
+    return _jax_wrapped_poisson_exact
+
+
+def _function_geometric_exact_named():
+    
+    def _jax_wrapped_geometric_exact(key, prob, param):
+        return random.geometric(key=key, p=prob, dtype=jnp.int64)
+    
+    return _jax_wrapped_geometric_exact
+
+
 class JaxRDDLCompiler:
     '''Compiles a RDDL AST representation into an equivalent JAX representation.
     All operations are identical to their numpy equivalents.
@@ -211,12 +221,12 @@ class JaxRDDLCompiler:
     }
     
     EXACT_RDDL_TO_JAX_IF = _function_if_exact_named()
-    
     EXACT_RDDL_TO_JAX_SWITCH = _function_switch_exact_named()
     
     EXACT_RDDL_TO_JAX_BERNOULLI = _function_bernoulli_exact_named()
-    
     EXACT_RDDL_TO_JAX_DISCRETE = _function_discrete_exact_named()
+    EXACT_RDDL_TO_JAX_POISSON = _function_poisson_exact_named()
+    EXACT_RDDL_TO_JAX_GEOMETRIC = _function_geometric_exact_named()
 
     def __init__(self, rddl: RDDLLiftedModel,
                  allow_synchronous_state: bool=True,
@@ -290,6 +300,8 @@ class JaxRDDLCompiler:
         self.SWITCH_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_SWITCH
         self.BERNOULLI_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_BERNOULLI
         self.DISCRETE_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_DISCRETE
+        self.POISSON_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_POISSON
+        self.GEOMETRIC_HELPER = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_GEOMETRIC
     
     # ===========================================================================
     # main compilation subroutines
@@ -997,13 +1009,14 @@ class JaxRDDLCompiler:
             jax_op, jax_param = self._unwrap(negative_op, expr.id, info)
             return self._jax_unary(jax_expr, jax_op, jax_param, at_least_int=True)
                     
-        elif n == 2:
-            lhs, rhs = args
-            jax_lhs = self._jax(lhs, info)
-            jax_rhs = self._jax(rhs, info)
+        elif n == 2 or (n >= 2 and op in {'*', '+'}):
+            jax_exprs = [self._jax(arg, info) for arg in args]
             jax_op, jax_param = self._unwrap(valid_ops[op], expr.id, info)
-            return self._jax_binary(
-                jax_lhs, jax_rhs, jax_op, jax_param, at_least_int=True)
+            result = jax_exprs[0]
+            for jax_rhs in jax_exprs[1:]:
+                result = self._jax_binary(
+                    result, jax_rhs, jax_op, jax_param, at_least_int=True)
+            return result
         
         JaxRDDLCompiler._check_num_args(expr, 2)
     
@@ -1047,13 +1060,14 @@ class JaxRDDLCompiler:
             jax_op, jax_param = self._unwrap(logical_not_op, expr.id, info)
             return self._jax_unary(jax_expr, jax_op, jax_param, check_dtype=bool)
         
-        elif n == 2:
-            lhs, rhs = args
-            jax_lhs = self._jax(lhs, info)
-            jax_rhs = self._jax(rhs, info)
+        elif n == 2 or (n >= 2 and op in {'^', '&', '|'}):
+            jax_exprs = [self._jax(arg, info) for arg in args]
             jax_op, jax_param = self._unwrap(valid_ops[op], expr.id, info)
-            return self._jax_binary(
-                jax_lhs, jax_rhs, jax_op, jax_param, check_dtype=bool)
+            result = jax_exprs[0]
+            for jax_rhs in jax_exprs[1:]:
+                result = self._jax_binary(
+                    result, jax_rhs, jax_op, jax_param, check_dtype=bool)
+            return result
         
         JaxRDDLCompiler._check_num_args(expr, 2)
     
@@ -1166,16 +1180,17 @@ class JaxRDDLCompiler:
         return _jax_wrapped_if_then_else
     
     def _jax_switch(self, expr, info):
+        pred, *_ = expr.args
              
-        # if expression is non-fluent, always use the exact operation
-        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(expr):
+        # if predicate is non-fluent, always use the exact operation
+        # case conditions are currently only literals so they are non-fluent
+        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(pred):
             switch_op = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_SWITCH
         else:
             switch_op = self.SWITCH_HELPER
         jax_switch, jax_param = self._unwrap(switch_op, expr.id, info)
         
         # recursively compile predicate
-        pred, *_ = expr.args  
         jax_pred = self._jax(pred, info)
         
         # recursively compile cases
@@ -1427,15 +1442,24 @@ class JaxRDDLCompiler:
     def _jax_poisson(self, expr, info):
         ERR = JaxRDDLCompiler.ERROR_CODES['INVALID_PARAM_POISSON']
         JaxRDDLCompiler._check_num_args(expr, 1)
-        
         arg_rate, = expr.args
+        
+        # if rate is non-fluent, always use the exact operation
+        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(arg_rate):
+            poisson_op = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_POISSON
+        else:
+            poisson_op = self.POISSON_HELPER
+        jax_poisson, jax_param = self._unwrap(poisson_op, expr.id, info)
+        
+        # recursively compile arguments
         jax_rate = self._jax(arg_rate, info)
         
         # uses the implicit JAX subroutine
         def _jax_wrapped_distribution_poisson(x, params, key):
             rate, key, err = jax_rate(x, params, key)
             key, subkey = random.split(key)
-            sample = random.poisson(key=subkey, lam=rate, dtype=self.INT)
+            param = params.get(jax_param, None)
+            sample = jax_poisson(subkey, rate, param).astype(self.INT)
             out_of_bounds = jnp.logical_not(jnp.all(rate >= 0))
             err |= (out_of_bounds * ERR)
             return sample, key, err
@@ -1536,33 +1560,25 @@ class JaxRDDLCompiler:
         ERR = JaxRDDLCompiler.ERROR_CODES['INVALID_PARAM_GEOMETRIC']
         JaxRDDLCompiler._check_num_args(expr, 1)        
         arg_prob, = expr.args
+        
+        # if prob is non-fluent, always use the exact operation
+        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(arg_prob):
+            geom_op = JaxRDDLCompiler.EXACT_RDDL_TO_JAX_GEOMETRIC
+        else:
+            geom_op = self.GEOMETRIC_HELPER
+        jax_geom, jax_param = self._unwrap(geom_op, expr.id, info)
+        
+        # recursively compile arguments        
         jax_prob = self._jax(arg_prob, info)
         
-        if self.compile_non_fluent_exact and not self.traced.cached_is_fluent(arg_prob):            
-            
-            # prob is non-fluent: do not reparameterize
-            def _jax_wrapped_distribution_geometric(x, params, key):
-                prob, key, err = jax_prob(x, params, key)
-                key, subkey = random.split(key)
-                sample = random.geometric(key=subkey, p=prob, dtype=self.INT)
-                out_of_bounds = jnp.logical_not(jnp.all((prob >= 0) & (prob <= 1)))
-                err |= (out_of_bounds * ERR)
-                return sample, key, err
-        
-        else:    
-            floor_op, jax_param = self._unwrap(
-                self.KNOWN_UNARY['floor'], expr.id, info)
-            
-            # reparameterization trick Geom(p) = floor(ln(U(0, 1)) / ln(p)) + 1
-            def _jax_wrapped_distribution_geometric(x, params, key):
-                prob, key, err = jax_prob(x, params, key)
-                key, subkey = random.split(key)
-                U = random.uniform(key=subkey, shape=jnp.shape(prob), dtype=self.REAL)
-                param = params.get(jax_param, None)
-                sample = floor_op(jnp.log(U) / jnp.log(1.0 - prob), param) + 1
-                out_of_bounds = jnp.logical_not(jnp.all((prob >= 0) & (prob <= 1)))
-                err |= (out_of_bounds * ERR)
-                return sample, key, err
+        def _jax_wrapped_distribution_geometric(x, params, key):
+            prob, key, err = jax_prob(x, params, key)
+            key, subkey = random.split(key)
+            param = params.get(jax_param, None)
+            sample = jax_geom(subkey, prob, param).astype(self.INT)
+            out_of_bounds = jnp.logical_not(jnp.all((prob >= 0) & (prob <= 1)))
+            err |= (out_of_bounds * ERR)
+            return sample, key, err
         
         return _jax_wrapped_distribution_geometric
     
@@ -1771,7 +1787,10 @@ class JaxRDDLCompiler:
             # dispatch to sampling subroutine
             key, subkey = random.split(key)
             param = params.get(jax_param, None)
-            sample, out_of_bounds = jax_discrete(subkey, prob, param)
+            sample = jax_discrete(subkey, prob, param)
+            out_of_bounds = jnp.logical_not(jnp.logical_and(
+                jnp.all(prob >= 0),
+                jnp.allclose(jnp.sum(prob, axis=-1), 1.0)))
             error |= (out_of_bounds * ERR)
             return sample, key, error
         
@@ -1804,7 +1823,10 @@ class JaxRDDLCompiler:
             # dispatch to sampling subroutine
             key, subkey = random.split(key)
             param = params.get(jax_param, None)
-            sample, out_of_bounds = jax_discrete(subkey, prob, param)
+            sample = jax_discrete(subkey, prob, param)
+            out_of_bounds = jnp.logical_not(jnp.logical_and(
+                jnp.all(prob >= 0),
+                jnp.allclose(jnp.sum(prob, axis=-1), 1.0)))
             error |= (out_of_bounds * ERR)
             return sample, key, error
         

pyRDDLGym_jax/core/logic.py

@@ -1,11 +1,19 @@
+from typing import Optional, Set
+
 import jax
 import jax.numpy as jnp
 import jax.random as random
-from typing import Optional, Set
 
 from pyRDDLGym.core.debug.exception import raise_warning
 
 
+# ===========================================================================
+# LOGICAL COMPLEMENT
+# - abstract class
+# - standard complement
+#
+# ===========================================================================
+
 class Complement:
     '''Base class for approximate logical complement operations.'''
     
@@ -20,6 +28,13 @@ class StandardComplement(Complement):
         return 1.0 - x
 
 
+# ===========================================================================
+# RELATIONAL OPERATIONS
+# - abstract class
+# - sigmoid comparison
+#
+# ===========================================================================
+
 class Comparison:
     '''Base class for approximate comparison operations.'''
     
@@ -44,7 +59,17 @@ class SigmoidComparison(Comparison):
     
     def equal(self, x, y, param):
         return 1.0 - jnp.square(jnp.tanh(param * (y - x)))
-        
+ 
+        
+# ===========================================================================
+# TNORMS
+# - abstract tnorm
+# - product tnorm
+# - Godel tnorm
+# - Lukasiewicz tnorm
+# - Yager(p) tnorm
+#
+# ===========================================================================
 
 class TNorm:
     '''Base class for fuzzy differentiable t-norms.'''
@@ -86,8 +111,133 @@ class LukasiewiczTNorm(TNorm):
     
     def norms(self, x, axis):
         return jax.nn.relu(jnp.sum(x - 1.0, axis=axis) + 1.0)
+
+
+class YagerTNorm(TNorm):
+    '''Yager t-norm given by the expression 
+    (x, y) -> max(1 - ((1 - x)^p + (1 - y)^p)^(1/p)).'''
+    
+    def __init__(self, p=2.0):
+        self.p = p
+    
+    def norm(self, x, y):
+        base_x = jax.nn.relu(1.0 - x)
+        base_y = jax.nn.relu(1.0 - y)
+        arg = jnp.power(base_x ** self.p + base_y ** self.p, 1.0 / self.p)
+        return jax.nn.relu(1.0 - arg)
+    
+    def norms(self, x, axis):
+        base = jax.nn.relu(1.0 - x)
+        arg = jnp.power(jnp.sum(base ** self.p, axis=axis), 1.0 / self.p)
+        return jax.nn.relu(1.0 - arg)
+        
+
+# ===========================================================================
+# RANDOM SAMPLING
+# - abstract sampler
+# - Gumbel-softmax sampler
+# - determinization
+#
+# ===========================================================================
+
+class RandomSampling:
+    '''An abstract class that describes how discrete and non-reparameterizable 
+    random variables are sampled.'''
+    
+    def discrete(self, logic):
+        raise NotImplementedError
+    
+    def bernoulli(self, logic):
+        jax_discrete, jax_param = self.discrete(logic)
+        
+        def _jax_wrapped_calc_bernoulli_approx(key, prob, param):
+            prob = jnp.stack([1.0 - prob, prob], axis=-1)
+            sample = jax_discrete(key, prob, param)
+            return sample
+        
+        return _jax_wrapped_calc_bernoulli_approx, jax_param
+    
+    def poisson(self, logic):
+        
+        def _jax_wrapped_calc_poisson_exact(key, rate, param):
+            return random.poisson(key=key, lam=rate, dtype=logic.INT)
+        
+        return _jax_wrapped_calc_poisson_exact, None
+    
+    def geometric(self, logic):
+        if logic.verbose:
+            raise_warning('Using the replacement rule: '
+                          'Geometric(p) --> floor(log(U) / log(1 - p)) + 1')
+            
+        jax_floor, jax_param = logic.floor()
+            
+        def _jax_wrapped_calc_geometric_approx(key, prob, param):
+            U = random.uniform(key=key, shape=jnp.shape(prob), dtype=logic.REAL)
+            sample = jax_floor(jnp.log(U) / jnp.log(1.0 - prob), param) + 1
+            return sample
+        
+        return _jax_wrapped_calc_geometric_approx, jax_param
+
+
+class GumbelSoftmax(RandomSampling):
+    '''Random sampling of discrete variables using Gumbel-softmax trick.'''
     
+    def discrete(self, logic):
+        if logic.verbose:
+            raise_warning('Using the replacement rule: '
+                          'Discrete(p) --> Gumbel-softmax(p)')
+        
+        jax_argmax, jax_param = logic.argmax()
+        
+        def _jax_wrapped_calc_discrete_gumbel_softmax(key, prob, param):
+            Gumbel01 = random.gumbel(key=key, shape=prob.shape, dtype=logic.REAL)
+            sample = Gumbel01 + jnp.log(prob + logic.eps)
+            sample = jax_argmax(sample, axis=-1, param=param)
+            return sample
+        
+        return _jax_wrapped_calc_discrete_gumbel_softmax, jax_param
+    
+
+class Determinization(RandomSampling):
+    '''Random sampling of variables using their deterministic mean estimate.'''
     
+    def discrete(self, logic):
+        if logic.verbose:
+            raise_warning('Using the replacement rule: '
+                          'Discrete(p) --> sum(i * p[i])')
+        
+        def _jax_wrapped_calc_discrete_determinized(key, prob, param):
+            literals = FuzzyLogic.enumerate_literals(prob.shape, axis=-1)
+            sample = jnp.sum(literals * prob, axis=-1)
+            return sample
+        
+        return _jax_wrapped_calc_discrete_determinized, None
+    
+    def poisson(self, logic):
+        if logic.verbose:
+            raise_warning('Using the replacement rule: Poisson(rate) --> rate')
+            
+        def _jax_wrapped_calc_poisson_determinized(key, rate, param):
+            return rate
+        
+        return _jax_wrapped_calc_poisson_determinized, None
+    
+    def geometric(self, logic):
+        if logic.verbose:
+            raise_warning('Using the replacement rule: Geometric(p) --> 1 / p')
+            
+        def _jax_wrapped_calc_geometric_determinized(key, prob, param):
+            sample = 1.0 / prob
+            return sample
+        
+        return _jax_wrapped_calc_geometric_determinized, None
+
+
+# ===========================================================================
+# FUZZY LOGIC
+#
+# ===========================================================================
+
 class FuzzyLogic:
     '''A class representing fuzzy logic in JAX.
     
@@ -98,9 +248,10 @@ class FuzzyLogic:
     def __init__(self, tnorm: TNorm=ProductTNorm(),
                  complement: Complement=StandardComplement(),
                  comparison: Comparison=SigmoidComparison(),
+                 sampling: RandomSampling=GumbelSoftmax(),
                  weight: float=10.0,
                  debias: Optional[Set[str]]=None,
-                 eps: float=1e-10,
+                 eps: float=1e-15,
                  verbose: bool=False,
                  use64bit: bool=False) -> None:
         '''Creates a new fuzzy logic in Jax.
@@ -108,8 +259,8 @@ class FuzzyLogic:
         :param tnorm: fuzzy operator for logical AND
         :param complement: fuzzy operator for logical NOT
         :param comparison: fuzzy operator for comparisons (>, >=, <, ==, ~=, ...)
+        :param sampling: random sampling of non-reparameterizable distributions
         :param weight: a sharpness parameter for sigmoid and softmax activations
-        :param error: an error parameter (e.g. floor) (smaller means better accuracy)
         :param debias: which functions to de-bias approximate on forward pass
         :param eps: small positive float to mitigate underflow
         :param verbose: whether to dump replacements and other info to console
@@ -118,6 +269,7 @@ class FuzzyLogic:
         self.tnorm = tnorm
         self.complement = complement
         self.comparison = comparison
+        self.sampling = sampling
         self.weight = float(weight)
         if debias is None:
             debias = set()
@@ -142,10 +294,11 @@ class FuzzyLogic:
               f'    tnorm         ={type(self.tnorm).__name__}\n'
               f'    complement    ={type(self.complement).__name__}\n'
               f'    comparison    ={type(self.comparison).__name__}\n'
+              f'    sampling      ={type(self.sampling).__name__}\n'
               f'    sigmoid_weight={self.weight}\n'
               f'    cpfs_to_debias={self.debias}\n'
               f'    underflow_tol ={self.eps}\n'
-              f'    use64bit      ={self.use64bit}')
+              f'    use_64_bit    ={self.use64bit}')
         
     # ===========================================================================
     # logical operators
@@ -419,7 +572,7 @@ class FuzzyLogic:
     # ===========================================================================
      
     @staticmethod
-    def _literals(shape, axis):
+    def enumerate_literals(shape, axis):
         literals = jnp.arange(shape[axis])
         literals = literals[(...,) + (jnp.newaxis,) * (len(shape) - 1)]
         literals = jnp.moveaxis(literals, source=0, destination=axis)
@@ -434,7 +587,7 @@ class FuzzyLogic:
         debias = 'argmax' in self.debias
         
         def _jax_wrapped_calc_argmax_approx(x, axis, param):
-            literals = FuzzyLogic._literals(x.shape, axis=axis)
+            literals = FuzzyLogic.enumerate_literals(x.shape, axis=axis)
             soft_max = jax.nn.softmax(param * x, axis=axis)
             sample = jnp.sum(literals * soft_max, axis=axis)
             if debias:
@@ -468,7 +621,7 @@ class FuzzyLogic:
         def _jax_wrapped_calc_if_approx(c, a, b, param):
             sample = c * a + (1.0 - c) * b
             if debias:
-                hard_sample = jnp.select([c, ~c], [a, b])
+                hard_sample = jnp.where(c > 0.5, a, b)
                 sample += jax.lax.stop_gradient(hard_sample - sample)
             return sample
         
@@ -483,7 +636,7 @@ class FuzzyLogic:
         debias = 'switch' in self.debias
         
         def _jax_wrapped_calc_switch_approx(pred, cases, param):
-            literals = FuzzyLogic._literals(cases.shape, axis=0)
+            literals = FuzzyLogic.enumerate_literals(cases.shape, axis=0)
             pred = jnp.broadcast_to(pred[jnp.newaxis, ...], shape=cases.shape)
             proximity = -jnp.abs(pred - literals)
             soft_case = jax.nn.softmax(param * proximity, axis=0)
@@ -502,44 +655,24 @@ class FuzzyLogic:
     # random variables
     # ===========================================================================
      
-    def _gumbel_softmax(self, key, prob):
-        Gumbel01 = random.gumbel(key=key, shape=prob.shape, dtype=self.REAL)
-        sample = Gumbel01 + jnp.log(prob + self.eps)
-        return sample
-        
+    def discrete(self):
+        return self.sampling.discrete(self)
+    
     def bernoulli(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'Bernoulli(p) --> Gumbel-softmax(p)')
-        
-        jax_gs = self._gumbel_softmax
-        jax_argmax, jax_param = self.argmax()
-        
-        def _jax_wrapped_calc_bernoulli_approx(key, prob, param):
-            prob = jnp.stack([1.0 - prob, prob], axis=-1)
-            sample = jax_gs(key, prob)
-            sample = jax_argmax(sample, axis=-1, param=param)
-            return sample
-        
-        return _jax_wrapped_calc_bernoulli_approx, jax_param
+        return self.sampling.bernoulli(self)
     
-    def discrete(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'Discrete(p) --> Gumbel-softmax(p)')
-        
-        jax_gs = self._gumbel_softmax
-        jax_argmax, jax_param = self.argmax()
-        
-        def _jax_wrapped_calc_discrete_approx(key, prob, param):
-            sample = jax_gs(key, prob) 
-            sample = jax_argmax(sample, axis=-1, param=param)
-            return sample
-        
-        return _jax_wrapped_calc_discrete_approx, jax_param
-        
+    def poisson(self):
+        return self.sampling.poisson(self)
+    
+    def geometric(self):
+        return self.sampling.geometric(self)
+
 
+# ===========================================================================
 # UNIT TESTS
+#
+# ===========================================================================
+
 logic = FuzzyLogic()
 w = 100.0
 
@@ -598,13 +731,14 @@ def _test_random():
     key = random.PRNGKey(42)
     _bernoulli, _ = logic.bernoulli()
     _discrete, _ = logic.discrete()
+    _geometric, _ = logic.geometric()
     
     def bern(n):
         prob = jnp.asarray([0.3] * n)
         sample = _bernoulli(key, prob, w)
         return sample
     
-    samples = bern(5000)
+    samples = bern(50000)
     print(jnp.mean(samples))
     
     def disc(n):
@@ -613,10 +747,18 @@ def _test_random():
         sample = _discrete(key, prob, w)
         return sample
         
-    samples = disc(5000)
+    samples = disc(50000)
     samples = jnp.round(samples)
     print([jnp.mean(samples == i) for i in range(3)])
-
+    
+    def geom(n):
+        prob = jnp.asarray([0.3] * n)
+        sample = _geometric(key, prob, w)
+        return sample
+    
+    samples = geom(50000)
+    print(jnp.mean(samples))
+    
 
 def _test_rounding():
     print('testing rounding')