pyRDDLGym-jax 0.5__py3-none-any.whl → 1.0__py3-none-any.whl

pyRDDLGym_jax/core/logic.py

@@ -4,7 +4,13 @@ import jax
 import jax.numpy as jnp
 import jax.random as random
 
-from pyRDDLGym.core.debug.exception import raise_warning
+
+def enumerate_literals(shape, axis, dtype=jnp.int32):
+    literals = jnp.arange(shape[axis], dtype=dtype)
+    literals = literals[(...,) + (jnp.newaxis,) * (len(shape) - 1)]
+    literals = jnp.moveaxis(literals, source=0, destination=axis)
+    literals = jnp.broadcast_to(literals, shape=shape)
+    return literals
 
 
 # ===========================================================================
@@ -17,36 +23,71 @@ from pyRDDLGym.core.debug.exception import raise_warning
 class Comparison:
     '''Base class for approximate comparison operations.'''
     
-    def greater_equal(self, x, y, param):
+    def greater_equal(self, id, init_params):
+        raise NotImplementedError
+    
+    def greater(self, id, init_params):
         raise NotImplementedError
     
-    def greater(self, x, y, param):
+    def equal(self, id, init_params):
         raise NotImplementedError
     
-    def equal(self, x, y, param):
+    def sgn(self, id, init_params):
         raise NotImplementedError
     
-    def sgn(self, x, param):
+    def argmax(self, id, init_params):
         raise NotImplementedError
     
 
 class SigmoidComparison(Comparison):
     '''Comparison operations approximated using sigmoid functions.'''
     
+    def __init__(self, weight: float=10.0):
+        self.weight = weight
+        
     # https://arxiv.org/abs/2110.05651
-    def greater_equal(self, x, y, param):
-        return jax.nn.sigmoid(param * (x - y))
-    
-    def greater(self, x, y, param):
-        return jax.nn.sigmoid(param * (x - y))
-    
-    def equal(self, x, y, param):
-        return 1.0 - jnp.square(jnp.tanh(param * (y - x)))
+    def greater_equal(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.weight
+        def _jax_wrapped_calc_greater_equal_approx(x, y, params):
+            gre_eq = jax.nn.sigmoid(params[id_] * (x - y))
+            return gre_eq, params
+        return _jax_wrapped_calc_greater_equal_approx
+    
+    def greater(self, id, init_params):
+        return self.greater_equal(id, init_params)
+    
+    def equal(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.weight
+        def _jax_wrapped_calc_equal_approx(x, y, params):
+            equal = 1.0 - jnp.square(jnp.tanh(params[id_] * (y - x)))
+            return equal, params
+        return _jax_wrapped_calc_equal_approx
+    
+    def sgn(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.weight
+        def _jax_wrapped_calc_sgn_approx(x, params):
+            sgn = jnp.tanh(params[id_] * x)
+            return sgn, params
+        return _jax_wrapped_calc_sgn_approx
     
-    def sgn(self, x, param):
-        return jnp.tanh(param * x)
- 
- 
+    # https://arxiv.org/abs/2110.05651
+    def argmax(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.weight
+        def _jax_wrapped_calc_argmax_approx(x, axis, params):
+            literals = enumerate_literals(x.shape, axis=axis)
+            softmax = jax.nn.softmax(params[id_] * x, axis=axis)
+            sample = jnp.sum(literals * softmax, axis=axis)
+            return sample, params
+        return _jax_wrapped_calc_argmax_approx
+    
+    def __str__(self) -> str:
+        return f'Sigmoid comparison with weight {self.weight}'
+
+
 # ===========================================================================
 # ROUNDING OPERATIONS
 # - abstract class
@@ -57,26 +98,44 @@ class SigmoidComparison(Comparison):
 class Rounding:
     '''Base class for approximate rounding operations.'''
     
-    def floor(self, x, param):
+    def floor(self, id, init_params):
         raise NotImplementedError
     
-    def round(self, x, param):
+    def round(self, id, init_params):
         raise NotImplementedError
 
 
 class SoftRounding(Rounding):
     '''Rounding operations approximated using soft operations.'''
     
+    def __init__(self, weight: float=10.0):
+        self.weight = weight
+        
     # https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/bijectors/Softfloor
-    def floor(self, x, param):
-        denom = jnp.tanh(param / 4.0)
-        return (jax.nn.sigmoid(param * (x - jnp.floor(x) - 1.0)) - 
-                jax.nn.sigmoid(-param / 2.0)) / denom + jnp.floor(x)
+    def floor(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.weight
+        def _jax_wrapped_calc_floor_approx(x, params):
+            param = params[id_]
+            denom = jnp.tanh(param / 4.0)
+            floor = (jax.nn.sigmoid(param * (x - jnp.floor(x) - 1.0)) - 
+                     jax.nn.sigmoid(-param / 2.0)) / denom + jnp.floor(x)
+            return floor, params
+        return _jax_wrapped_calc_floor_approx
     
     # https://arxiv.org/abs/2006.09952
-    def round(self, x, param):
-        m = jnp.floor(x) + 0.5
-        return m + 0.5 * jnp.tanh(param * (x - m)) / jnp.tanh(param / 2.0)    
+    def round(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.weight
+        def _jax_wrapped_calc_round_approx(x, params):
+            param = params[id_]
+            m = jnp.floor(x) + 0.5
+            rounded = m + 0.5 * jnp.tanh(param * (x - m)) / jnp.tanh(param / 2.0)
+            return rounded, params
+        return _jax_wrapped_calc_round_approx
+
+    def __str__(self) -> str:
+        return f'SoftFloor and SoftRound with weight {self.weight}'
 
 
 # ===========================================================================
@@ -89,7 +148,7 @@ class SoftRounding(Rounding):
 class Complement:
     '''Base class for approximate logical complement operations.'''
     
-    def __call__(self, x):
+    def __call__(self, id, init_params):
         raise NotImplementedError
 
 
@@ -97,9 +156,16 @@ class StandardComplement(Complement):
     '''The standard approximate logical complement given by x -> 1 - x.'''
     
     # https://www.sciencedirect.com/science/article/abs/pii/016501149190171L
-    def __call__(self, x):
-        return 1.0 - x
-
+    @staticmethod
+    def _jax_wrapped_calc_not_approx(x, params):
+        return 1.0 - x, params   
+    
+    def __call__(self, id, init_params):
+        return self._jax_wrapped_calc_not_approx
+    
+    def __str__(self) -> str:
+        return 'Standard complement'
+    
 
 # ===========================================================================
 # TNORMS
@@ -115,43 +181,78 @@ class StandardComplement(Complement):
 class TNorm:
     '''Base class for fuzzy differentiable t-norms.'''
     
-    def norm(self, x, y):
+    def norm(self, id, init_params):
         '''Elementwise t-norm of x and y.'''
         raise NotImplementedError
     
-    def norms(self, x, axis):
+    def norms(self, id, init_params):
         '''T-norm computed for tensor x along axis.'''
         raise NotImplementedError
-        
+    
 
 class ProductTNorm(TNorm):
     '''Product t-norm given by the expression (x, y) -> x * y.'''
     
-    def norm(self, x, y):
-        return x * y
+    @staticmethod
+    def _jax_wrapped_calc_and_approx(x, y, params):
+        return x * y, params    
+        
+    def norm(self, id, init_params):        
+        return self._jax_wrapped_calc_and_approx
     
-    def norms(self, x, axis):
-        return jnp.prod(x, axis=axis)
+    @staticmethod
+    def _jax_wrapped_calc_forall_approx(x, axis, params):
+        return jnp.prod(x, axis=axis), params   
+        
+    def norms(self, id, init_params):        
+        return self._jax_wrapped_calc_forall_approx
 
+    def __str__(self) -> str:
+        return 'Product t-norm'
+    
 
 class GodelTNorm(TNorm):
     '''Godel t-norm given by the expression (x, y) -> min(x, y).'''
     
-    def norm(self, x, y):
-        return jnp.minimum(x, y)
+    @staticmethod
+    def _jax_wrapped_calc_and_approx(x, y, params):
+        return jnp.minimum(x, y), params
+    
+    def norm(self, id, init_params):        
+        return self._jax_wrapped_calc_and_approx
+    
+    @staticmethod
+    def _jax_wrapped_calc_forall_approx(x, axis, params):
+        return jnp.min(x, axis=axis), params   
+        
+    def norms(self, id, init_params):        
+        return self._jax_wrapped_calc_forall_approx
+    
+    def __str__(self) -> str:
+        return 'Godel t-norm'
     
-    def norms(self, x, axis):
-        return jnp.min(x, axis=axis)
-
 
 class LukasiewiczTNorm(TNorm):
     '''Lukasiewicz t-norm given by the expression (x, y) -> max(x + y - 1, 0).'''
     
-    def norm(self, x, y):
-        return jax.nn.relu(x + y - 1.0)
+    @staticmethod
+    def _jax_wrapped_calc_and_approx(x, y, params):
+        land = jax.nn.relu(x + y - 1.0)
+        return land, params
+        
+    def norm(self, id, init_params):
+        return self._jax_wrapped_calc_and_approx
     
-    def norms(self, x, axis):
-        return jax.nn.relu(jnp.sum(x - 1.0, axis=axis) + 1.0)
+    @staticmethod
+    def _jax_wrapped_calc_forall_approx(x, axis, params):
+        forall = jax.nn.relu(jnp.sum(x - 1.0, axis=axis) + 1.0)
+        return forall, params
+        
+    def norms(self, id, init_params):
+        return self._jax_wrapped_calc_forall_approx
+    
+    def __str__(self) -> str:
+        return 'Lukasiewicz t-norm'
 
 
 class YagerTNorm(TNorm):
@@ -161,17 +262,30 @@ class YagerTNorm(TNorm):
     def __init__(self, p=2.0):
         self.p = float(p)
     
-    def norm(self, x, y):        
-        base = jax.nn.relu(1.0 - jnp.stack([x, y], axis=0))
-        arg = jnp.linalg.norm(base, ord=self.p, axis=0)
-        return jax.nn.relu(1.0 - arg)
+    def norm(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.p
+        def _jax_wrapped_calc_and_approx(x, y, params):
+            base = jax.nn.relu(1.0 - jnp.stack([x, y], axis=0))
+            arg = jnp.linalg.norm(base, ord=params[id_], axis=0)
+            land = jax.nn.relu(1.0 - arg)
+            return land, params
+        return _jax_wrapped_calc_and_approx
+    
+    def norms(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.p
+        def _jax_wrapped_calc_forall_approx(x, axis, params):
+            arg = jax.nn.relu(1.0 - x)
+            for ax in sorted(axis, reverse=True):
+                arg = jnp.linalg.norm(arg, ord=params[id_], axis=ax)
+            forall = jax.nn.relu(1.0 - arg)
+            return forall, params
+        return _jax_wrapped_calc_forall_approx
+    
+    def __str__(self) -> str:
+        return f'Yager({self.p}) t-norm'
     
-    def norms(self, x, axis):
-        arg = jax.nn.relu(1.0 - x)
-        for ax in sorted(axis, reverse=True):
-            arg = jnp.linalg.norm(arg, ord=self.p, axis=ax)
-        return jax.nn.relu(1.0 - arg)
-        
 
 # ===========================================================================
 # RANDOM SAMPLING
@@ -185,117 +299,443 @@ class RandomSampling:
     '''An abstract class that describes how discrete and non-reparameterizable 
     random variables are sampled.'''
     
-    def discrete(self, logic):
+    def discrete(self, id, init_params, logic):
         raise NotImplementedError
     
-    def bernoulli(self, logic):
-        jax_discrete, jax_param = self.discrete(logic)
-        
-        def _jax_wrapped_calc_bernoulli_approx(key, prob, param):
+    def bernoulli(self, id, init_params, logic):
+        discrete_approx = self.discrete(id, init_params, logic)        
+        def _jax_wrapped_calc_bernoulli_approx(key, prob, params):
             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
+            return discrete_approx(key, prob, params)        
+        return _jax_wrapped_calc_bernoulli_approx
     
-    def poisson(self, logic):
-        
-        def _jax_wrapped_calc_poisson_exact(key, rate, param):
-            return random.poisson(key=key, lam=rate, dtype=logic.INT)
+    @staticmethod
+    def _jax_wrapped_calc_poisson_exact(key, rate, params):
+        sample = random.poisson(key=key, lam=rate, dtype=logic.INT)
+        return sample, params
         
-        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):
+    def poisson(self, id, init_params, logic):
+        return self._jax_wrapped_calc_poisson_exact
+    
+    def geometric(self, id, init_params, logic):
+        approx_floor = logic.floor(id, init_params)
+        def _jax_wrapped_calc_geometric_approx(key, prob, params):
             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
+            floor, params = approx_floor(jnp.log(U) / jnp.log(1.0 - prob), params)
+            sample = floor + 1
+            return sample, params
+        return _jax_wrapped_calc_geometric_approx
 
 
 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()
-        
-        # https://arxiv.org/pdf/1611.01144
-        def _jax_wrapped_calc_discrete_gumbel_softmax(key, prob, param):
+    # https://arxiv.org/pdf/1611.01144
+    def discrete(self, id, init_params, logic):
+        argmax_approx = logic.argmax(id, init_params)
+        def _jax_wrapped_calc_discrete_gumbel_softmax(key, prob, params):
             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
+            return argmax_approx(sample, axis=-1, params=params)
+        return _jax_wrapped_calc_discrete_gumbel_softmax
+    
+    def __str__(self) -> str:
+        return 'Gumbel-Softmax'
     
 
 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
+    @staticmethod
+    def _jax_wrapped_calc_discrete_determinized(key, prob, params):
+        literals = enumerate_literals(prob.shape, axis=-1)
+        sample = jnp.sum(literals * prob, axis=-1)
+        return sample, params
         
-        return _jax_wrapped_calc_geometric_determinized, None
+    def discrete(self, id, init_params, logic):
+        return self._jax_wrapped_calc_discrete_determinized
+    
+    @staticmethod
+    def _jax_wrapped_calc_poisson_determinized(key, rate, params):
+        return rate, params       
 
+    def poisson(self, id, init_params, logic):
+        return self._jax_wrapped_calc_poisson_determinized
+    
+    @staticmethod
+    def _jax_wrapped_calc_geometric_determinized(key, prob, params):
+        sample = 1.0 / prob
+        return sample, params   
+    
+    def geometric(self, id, init_params, logic):
+        return self._jax_wrapped_calc_geometric_determinized
+    
+    def __str__(self) -> str:
+        return 'Deterministic'
+    
 
 # ===========================================================================
-# FUZZY LOGIC
+# CONTROL FLOW
+# - soft flow
 #
 # ===========================================================================
 
-class FuzzyLogic:
-    '''A class representing fuzzy logic in JAX.
+class ControlFlow:
+    '''A base class for control flow, including if and switch statements.'''
+    
+    def if_then_else(self, id, init_params):
+        raise NotImplementedError
     
-    Functionality can be customized by either providing a tnorm as parameters, 
-    or by overriding its methods.
-    '''
+    def switch(self, id, init_params):
+        raise NotImplementedError
+
+
+class SoftControlFlow(ControlFlow):
+    '''Soft control flow using a probabilistic interpretation.'''
+    
+    def __init__(self, weight: float=10.0) -> None:
+        self.weight = weight
+        
+    @staticmethod
+    def _jax_wrapped_calc_if_then_else_soft(c, a, b, params):
+        sample = c * a + (1.0 - c) * b
+        return sample, params
+    
+    def if_then_else(self, id, init_params):
+        return self._jax_wrapped_calc_if_then_else_soft
+    
+    def switch(self, id, init_params):
+        id_ = str(id)
+        init_params[id_] = self.weight
+        def _jax_wrapped_calc_switch_soft(pred, cases, params):
+            literals = enumerate_literals(cases.shape, axis=0)
+            pred = jnp.broadcast_to(pred[jnp.newaxis, ...], shape=cases.shape)
+            proximity = -jnp.square(pred - literals)
+            softcase = jax.nn.softmax(params[id_] * proximity, axis=0)
+            sample = jnp.sum(cases * softcase, axis=0)
+            return sample, params
+        return _jax_wrapped_calc_switch_soft
+    
+    def __str__(self) -> str:
+        return f'Soft control flow with weight {self.weight}'
+    
+    
+# ===========================================================================
+# LOGIC
+# - exact logic
+# - fuzzy logic
+#
+# ===========================================================================
+
+
+class Logic:
+    '''A base class for representing logic computations in JAX.'''
+    
+    def __init__(self, use64bit: bool=False) -> None:
+        self.set_use64bit(use64bit)
+    
+    def summarize_hyperparameters(self) -> None:
+        print(f'model relaxation:\n'
+              f'    use_64_bit    ={self.use64bit}')
+    
+    def set_use64bit(self, use64bit: bool) -> None:
+        '''Toggles whether or not the JAX system will use 64 bit precision.'''
+        self.use64bit = use64bit
+        if use64bit:
+            self.REAL = jnp.float64
+            self.INT = jnp.int64
+            jax.config.update('jax_enable_x64', True)
+        else:
+            self.REAL = jnp.float32
+            self.INT = jnp.int32
+            jax.config.update('jax_enable_x64', False)
+    
+    # ===========================================================================
+    # logical operators
+    # ===========================================================================
+     
+    def logical_and(self, id, init_params):
+        raise NotImplementedError
+    
+    def logical_not(self, id, init_params):
+        raise NotImplementedError
+    
+    def logical_or(self, id, init_params):
+        raise NotImplementedError
+    
+    def xor(self, id, init_params):
+        raise NotImplementedError
+    
+    def implies(self, id, init_params):
+        raise NotImplementedError
+    
+    def equiv(self, id, init_params):
+        raise NotImplementedError
+    
+    def forall(self, id, init_params):
+        raise NotImplementedError
+    
+    def exists(self, id, init_params):    
+        raise NotImplementedError 
+    
+    # ===========================================================================
+    # comparison operators
+    # ===========================================================================
+    
+    def greater_equal(self, id, init_params):
+        raise NotImplementedError 
+    
+    def greater(self, id, init_params):
+        raise NotImplementedError 
+    
+    def less_equal(self, id, init_params):
+        raise NotImplementedError 
+    
+    def less(self, id, init_params):
+        raise NotImplementedError 
+    
+    def equal(self, id, init_params):
+        raise NotImplementedError 
+    
+    def not_equal(self, id, init_params):
+        raise NotImplementedError 
+    
+    # ===========================================================================
+    # special functions
+    # ===========================================================================
+     
+    def sgn(self, id, init_params):
+        raise NotImplementedError 
+    
+    def floor(self, id, init_params):
+        raise NotImplementedError 
+    
+    def round(self, id, init_params):
+        raise NotImplementedError 
+    
+    def ceil(self, id, init_params):
+        raise NotImplementedError 
+    
+    def div(self, id, init_params):
+        raise NotImplementedError 
+    
+    def mod(self, id, init_params):
+        raise NotImplementedError 
+    
+    def sqrt(self, id, init_params):
+        raise NotImplementedError 
+    
+    # ===========================================================================
+    # indexing
+    # ===========================================================================
+     
+    def argmax(self, id, init_params):   
+        raise NotImplementedError  
+    
+    def argmin(self, id, init_params):   
+        raise NotImplementedError     
+    
+    # ===========================================================================
+    # control flow
+    # ===========================================================================
+     
+    def control_if(self, id, init_params):
+        raise NotImplementedError
+        
+    def control_switch(self, id, init_params):
+        raise NotImplementedError
+    
+    # ===========================================================================
+    # random variables
+    # ===========================================================================
+     
+    def discrete(self, id, init_params):
+        raise NotImplementedError
+    
+    def bernoulli(self, id, init_params):
+        raise NotImplementedError
+    
+    def poisson(self, id, init_params):
+        raise NotImplementedError
+    
+    def geometric(self, id, init_params):
+        raise NotImplementedError
+
+
+class ExactLogic(Logic):
+    '''A class representing exact logic in JAX.'''
+    
+    @staticmethod
+    def exact_unary_function(op):
+        def _jax_wrapped_calc_unary_function_exact(x, params):
+            return op(x), params
+        return _jax_wrapped_calc_unary_function_exact
+    
+    @staticmethod
+    def exact_binary_function(op):
+        def _jax_wrapped_calc_binary_function_exact(x, y, params):
+            return op(x, y), params
+        return _jax_wrapped_calc_binary_function_exact
+    
+    @staticmethod
+    def exact_aggregation(op):        
+        def _jax_wrapped_calc_aggregation_exact(x, axis, params):
+            return op(x, axis=axis), params   
+        return _jax_wrapped_calc_aggregation_exact
+
+    # ===========================================================================
+    # logical operators
+    # ===========================================================================
+     
+    def logical_and(self, id, init_params):
+        return self.exact_binary_function(jnp.logical_and)
+    
+    def logical_not(self, id, init_params):
+        return self.exact_unary_function(jnp.logical_not)
+    
+    def logical_or(self, id, init_params):
+        return self.exact_binary_function(jnp.logical_or)
+    
+    def xor(self, id, init_params):
+        return self.exact_binary_function(jnp.logical_xor)
+    
+    @staticmethod
+    def exact_binary_implies(x, y, params):
+        return jnp.logical_or(jnp.logical_not(x), y), params     
+    
+    def implies(self, id, init_params):
+        return self.exact_binary_implies
+    
+    def equiv(self, id, init_params):
+        return self.exact_binary_function(jnp.equal)
+    
+    def forall(self, id, init_params):
+        return self.exact_aggregation(jnp.all)
+    
+    def exists(self, id, init_params):
+        return self.exact_aggregation(jnp.any)
+    
+    # ===========================================================================
+    # comparison operators
+    # ===========================================================================
+    
+    def greater_equal(self, id, init_params):
+        return self.exact_binary_function(jnp.greater_equal)
+    
+    def greater(self, id, init_params):
+        return self.exact_binary_function(jnp.greater)
+    
+    def less_equal(self, id, init_params):
+        return self.exact_binary_function(jnp.less_equal)
+    
+    def less(self, id, init_params):
+        return self.exact_binary_function(jnp.less)
+    
+    def equal(self, id, init_params):
+        return self.exact_binary_function(jnp.equal)
+    
+    def not_equal(self, id, init_params):
+        return self.exact_binary_function(jnp.not_equal)
+    
+    # ===========================================================================
+    # special functions
+    # ===========================================================================
+     
+    def sgn(self, id, init_params):
+        return self.exact_unary_function(jnp.sign)
+    
+    def floor(self, id, init_params):
+        return self.exact_unary_function(jnp.floor)
+    
+    def round(self, id, init_params):
+        return self.exact_unary_function(jnp.round)
+    
+    def ceil(self, id, init_params):
+        return self.exact_unary_function(jnp.ceil)
+    
+    def div(self, id, init_params):
+        return self.exact_binary_function(jnp.floor_divide)
+    
+    def mod(self, id, init_params):
+        return self.exact_binary_function(jnp.mod)
+    
+    def sqrt(self, id, init_params):
+        return self.exact_unary_function(jnp.sqrt)
+    
+    # ===========================================================================
+    # indexing
+    # ===========================================================================
+     
+    def argmax(self, id, init_params):   
+        return self.exact_aggregation(jnp.argmax)
+    
+    def argmin(self, id, init_params):   
+        return self.exact_aggregation(jnp.argmin)
+    
+    # ===========================================================================
+    # control flow
+    # ===========================================================================
+    
+    @staticmethod
+    def exact_if_then_else(c, a, b, params):
+        return jnp.where(c > 0.5, a, b), params
+         
+    def control_if(self, id, init_params):
+        return self.exact_if_then_else
+    
+    @staticmethod
+    def exact_switch(pred, cases, params):
+        pred = pred[jnp.newaxis, ...]
+        sample = jnp.take_along_axis(cases, pred, axis=0)
+        assert sample.shape[0] == 1
+        return sample[0, ...], params
+    
+    def control_switch(self, id, init_params):
+        return self.exact_switch
+    
+    # ===========================================================================
+    # random variables
+    # ===========================================================================
+    
+    @staticmethod
+    def exact_discrete(key, prob, params):
+        return random.categorical(key=key, logits=jnp.log(prob), axis=-1), params  
+    
+    def discrete(self, id, init_params):
+        return self.exact_discrete
+    
+    @staticmethod
+    def exact_bernoulli(key, prob, params):
+        return random.bernoulli(key, prob), params
+    
+    def bernoulli(self, id, init_params):
+        return self.exact_bernoulli
+    
+    @staticmethod
+    def exact_poisson(key, rate, params):
+        return random.poisson(key=key, lam=rate), params   
+    
+    def poisson(self, id, init_params):
+        return self.exact_poisson
+    
+    @staticmethod
+    def exact_geometric(key, prob, params):
+        return random.geometric(key=key, p=prob), params
+    
+    def geometric(self, id, init_params):
+        return self.exact_geometric
+    
+    
+class FuzzyLogic(Logic):
+    '''A class representing fuzzy logic in JAX.'''
     
     def __init__(self, tnorm: TNorm=ProductTNorm(),
                  complement: Complement=StandardComplement(),
                  comparison: Comparison=SigmoidComparison(),
                  sampling: RandomSampling=GumbelSoftmax(),
                  rounding: Rounding=SoftRounding(),
-                 weight: float=10.0,
-                 debias: Optional[Set[str]]=None,
+                 control: ControlFlow=SoftControlFlow(),
                  eps: float=1e-15,
-                 verbose: bool=False,
                  use64bit: bool=False) -> None:
         '''Creates a new fuzzy logic in Jax.
         
@@ -304,428 +744,212 @@ class FuzzyLogic:
         :param comparison: fuzzy operator for comparisons (>, >=, <, ==, ~=, ...)
         :param sampling: random sampling of non-reparameterizable distributions
         :param rounding: rounding floating values to integers
-        :param weight: a sharpness parameter for sigmoid and softmax activations
-        :param debias: which functions to de-bias approximate on forward pass
+        :param control: if and switch control structures
         :param eps: small positive float to mitigate underflow
-        :param verbose: whether to dump replacements and other info to console
         :param use64bit: whether to perform arithmetic in 64 bit
         '''
+        super().__init__(use64bit=use64bit)
         self.tnorm = tnorm
         self.complement = complement
         self.comparison = comparison
         self.sampling = sampling
         self.rounding = rounding
-        self.weight = float(weight)
-        if debias is None:
-            debias = set()
-        self.debias = debias
+        self.control = control
         self.eps = eps
-        self.verbose = verbose
-        self.set_use64bit(use64bit)
     
-    def set_use64bit(self, use64bit: bool) -> None:
-        self.use64bit = use64bit
-        if use64bit:
-            self.REAL = jnp.float64
-            self.INT = jnp.int64
-            jax.config.update('jax_enable_x64', True)
-        else:
-            self.REAL = jnp.float32
-            self.INT = jnp.int32
-            jax.config.update('jax_enable_x64', False)
-        
+    def __str__(self) -> str:
+        return (f'model relaxation:\n'
+                f'    tnorm         ={str(self.tnorm)}\n'
+                f'    complement    ={str(self.complement)}\n'
+                f'    comparison    ={str(self.comparison)}\n'
+                f'    sampling      ={str(self.sampling)}\n'
+                f'    rounding      ={str(self.rounding)}\n'
+                f'    control       ={str(self.control)}\n'
+                f'    underflow_tol ={self.eps}\n'
+                f'    use_64_bit    ={self.use64bit}')
+
     def summarize_hyperparameters(self) -> None:
-        print(f'model relaxation:\n'
-              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'    rounding      ={type(self.rounding).__name__}\n'
-              f'    sigmoid_weight={self.weight}\n'
-              f'    cpfs_to_debias={self.debias}\n'
-              f'    underflow_tol ={self.eps}\n'
-              f'    use_64_bit    ={self.use64bit}')
+        print(self.__str__())
         
     # ===========================================================================
     # logical operators
     # ===========================================================================
      
-    def logical_and(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: a ^ b --> tnorm(a, b).')
-        
-        _and = self.tnorm.norm
-        
-        def _jax_wrapped_calc_and_approx(a, b, param):
-            return _and(a, b)
-        
-        return _jax_wrapped_calc_and_approx, None
-    
-    def logical_not(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: ~a --> complement(a)')
-        
-        _not = self.complement
-        
-        def _jax_wrapped_calc_not_approx(x, param):
-            return _not(x)
-        
-        return _jax_wrapped_calc_not_approx, None
-    
-    def logical_or(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: a | b --> tconorm(a, b).')
-            
-        _not = self.complement
-        _and = self.tnorm.norm
-        
-        def _jax_wrapped_calc_or_approx(a, b, param):
-            return _not(_and(_not(a), _not(b)))
-        
-        return _jax_wrapped_calc_or_approx, None
+    def logical_and(self, id, init_params):
+        return self.tnorm.norm(id, init_params)
+    
+    def logical_not(self, id, init_params):
+        return self.complement(id, init_params)
+    
+    def logical_or(self, id, init_params):
+        _not1 = self.complement(f'{id}_~1', init_params)
+        _not2 = self.complement(f'{id}_~2', init_params)
+        _and = self.tnorm.norm(f'{id}_^', init_params)
+        _not = self.complement(f'{id}_~', init_params)
+        
+        def _jax_wrapped_calc_or_approx(x, y, params):
+            not_x, params = _not1(x, params)
+            not_y, params = _not2(y, params)
+            not_x_and_not_y, params = _and(not_x, not_y, params)
+            return _not(not_x_and_not_y, params)        
+        return _jax_wrapped_calc_or_approx
 
-    def xor(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'a ~ b --> (a | b) ^ (a ^ b).')
-        
-        _not = self.complement
-        _and = self.tnorm.norm
-        
-        def _jax_wrapped_calc_xor_approx(a, b, param):
-            _or = _not(_and(_not(a), _not(b)))
-            return _and(_or(a, b), _not(_and(a, b)))
-        
-        return _jax_wrapped_calc_xor_approx, None
-        
-    def implies(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: a => b --> ~a ^ b')
-        
-        _not = self.complement
-        _and = self.tnorm.norm
-        
-        def _jax_wrapped_calc_implies_approx(a, b, param):
-            return _not(_and(a, _not(b)))
-        
-        return _jax_wrapped_calc_implies_approx, None
-    
-    def equiv(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'a <=> b --> (a => b) ^ (b => a)')
-        
-        _not = self.complement
-        _and = self.tnorm.norm
-        
-        def _jax_wrapped_calc_equiv_approx(a, b, param):
-            atob = _not(_and(a, _not(b)))
-            btoa = _not(_and(b, _not(a)))
-            return _and(atob, btoa)
-        
-        return _jax_wrapped_calc_equiv_approx, None
-    
-    def forall(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'forall(a) --> a[1] ^ a[2] ^ ...')
-        
-        _forall = self.tnorm.norms
-        
-        def _jax_wrapped_calc_forall_approx(x, axis, param):
-            return _forall(x, axis=axis)
-        
-        return _jax_wrapped_calc_forall_approx, None
-    
-    def exists(self):
-        _not = self.complement
-        jax_forall, jax_param = self.forall()
-        
-        def _jax_wrapped_calc_exists_approx(x, axis, param):
-            return _not(jax_forall(_not(x), axis, param))
-        
-        return _jax_wrapped_calc_exists_approx, jax_param
+    def xor(self, id, init_params):
+        _not = self.complement(f'{id}_~', init_params)
+        _and1 = self.tnorm.norm(f'{id}_^1', init_params)
+        _and2 = self.tnorm.norm(f'{id}_^2', init_params)
+        _or = self.logical_or(f'{id}_|', init_params)
+        
+        def _jax_wrapped_calc_xor_approx(x, y, params):
+            x_and_y, params = _and1(x, y, params)
+            not_x_and_y, params = _not(x_and_y, params)
+            x_or_y, params = _or(x, y, params)
+            return _and2(x_or_y, not_x_and_y, params)        
+        return _jax_wrapped_calc_xor_approx
+        
+    def implies(self, id, init_params):
+        _not = self.complement(f'{id}_~', init_params)
+        _or = self.logical_or(f'{id}_|', init_params)
+        
+        def _jax_wrapped_calc_implies_approx(x, y, params):
+            not_x, params = _not(x, params)
+            return _or(not_x, y, params)        
+        return _jax_wrapped_calc_implies_approx
+       
+    def equiv(self, id, init_params):
+        _implies1 = self.implies(f'{id}_=>1', init_params)
+        _implies2 = self.implies(f'{id}_=>2', init_params)
+        _and = self.tnorm.norm(f'{id}_^', init_params)
+           
+        def _jax_wrapped_calc_equiv_approx(x, y, params):
+            x_implies_y, params = _implies1(x, y, params)
+            y_implies_x, params = _implies2(y, x, params)
+            return _and(x_implies_y, y_implies_x, params)        
+        return _jax_wrapped_calc_equiv_approx
+    
+    def forall(self, id, init_params):
+        return self.tnorm.norms(id, init_params)
+    
+    def exists(self, id, init_params):
+        _not1 = self.complement(f'{id}_~1', init_params)
+        _not2 = self.complement(f'{id}_~2', init_params)
+        _forall = self.forall(f'{id}_forall', init_params)
+        
+        def _jax_wrapped_calc_exists_approx(x, axis, params):
+            not_x, params = _not1(x, params)
+            forall_not_x, params = _forall(not_x, axis, params)
+            return _not2(forall_not_x, params)        
+        return _jax_wrapped_calc_exists_approx
     
     # ===========================================================================
     # comparison operators
     # ===========================================================================
-     
-    def greater_equal(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'a >= b --> comparison.greater_equal(a, b)')
-        
-        greater_equal_op = self.comparison.greater_equal
-        debias = 'greater_equal' in self.debias
-        
-        def _jax_wrapped_calc_geq_approx(a, b, param):
-            sample = greater_equal_op(a, b, param)
-            if debias:
-                hard_sample = jnp.greater_equal(a, b)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'greater_equal')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_geq_approx, new_param
-    
-    def greater(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'a > b --> comparison.greater(a, b)')
-        
-        greater_op = self.comparison.greater
-        debias = 'greater' in self.debias
-        
-        def _jax_wrapped_calc_gre_approx(a, b, param):
-            sample = greater_op(a, b, param)
-            if debias:
-                hard_sample = jnp.greater(a, b)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'greater')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_gre_approx, new_param
     
-    def less_equal(self):
-        jax_geq, jax_param = self.greater_equal()
-        
-        def _jax_wrapped_calc_leq_approx(a, b, param):
-            return jax_geq(-a, -b, param)
-        
-        return _jax_wrapped_calc_leq_approx, jax_param
+    def greater_equal(self, id, init_params):
+        return self.comparison.greater_equal(id, init_params)
     
-    def less(self):
-        jax_gre, jax_param = self.greater()
-
-        def _jax_wrapped_calc_less_approx(a, b, param):
-            return jax_gre(-a, -b, param)
-        
-        return _jax_wrapped_calc_less_approx, jax_param
-
-    def equal(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'a == b --> comparison.equal(a, b)')
-        
-        equal_op = self.comparison.equal
-        debias = 'equal' in self.debias
-        
-        def _jax_wrapped_calc_equal_approx(a, b, param):
-            sample = equal_op(a, b, param)
-            if debias:
-                hard_sample = jnp.equal(a, b)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'equal')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_equal_approx, new_param
+    def greater(self, id, init_params):
+        return self.comparison.greater(id, init_params)
     
-    def not_equal(self):
-        _not = self.complement
-        jax_eq, jax_param = self.equal()
-        
-        def _jax_wrapped_calc_neq_approx(a, b, param):
-            return _not(jax_eq(a, b, param))
+    def less_equal(self, id, init_params):
+        _greater_eq = self.greater_equal(id, init_params)        
+        def _jax_wrapped_calc_leq_approx(x, y, params):
+            return _greater_eq(-x, -y, params)        
+        return _jax_wrapped_calc_leq_approx
+    
+    def less(self, id, init_params):
+        _greater = self.greater(id, init_params)  
+        def _jax_wrapped_calc_less_approx(x, y, params):
+            return _greater(-x, -y, params)
+        return _jax_wrapped_calc_less_approx
 
-        return _jax_wrapped_calc_neq_approx, jax_param
+    def equal(self, id, init_params):
+        return self.comparison.equal(id, init_params)
+    
+    def not_equal(self, id, init_params):
+        _not = self.complement(f'{id}_~', init_params)
+        _equal = self.comparison.equal(f'{id}_==', init_params)
+        def _jax_wrapped_calc_neq_approx(x, y, params):
+            equal, params = _equal(x, y, params)
+            return _not(equal, params)
+        return _jax_wrapped_calc_neq_approx
         
     # ===========================================================================
     # special functions
     # ===========================================================================
-     
-    def sgn(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'sgn(x) --> comparison.sgn(x)')
-        
-        sgn_op = self.comparison.sgn
-        debias = 'sgn' in self.debias
-        
-        def _jax_wrapped_calc_sgn_approx(x, param):
-            sample = sgn_op(x, param)
-            if debias:
-                hard_sample = jnp.sign(x)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'sgn')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_sgn_approx, new_param
-    
-    def floor(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'floor(x) --> rounding.floor(x)')
-        
-        floor_op = self.rounding.floor
-        debias = 'floor' in self.debias
-        
-        def _jax_wrapped_calc_floor_approx(x, param):
-            sample = floor_op(x, param)
-            if debias:
-                hard_sample = jnp.floor(x)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'floor')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_floor_approx, new_param
-        
-    def round(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'round(x) --> rounding.round(x)')
-        
-        round_op = self.rounding.round
-        debias = 'round' in self.debias
-        
-        def _jax_wrapped_calc_round_approx(x, param):
-            sample = round_op(x, param)
-            if debias:
-                hard_sample = jnp.round(x)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'round')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_round_approx, new_param
     
-    def ceil(self):
-        jax_floor, jax_param = self.floor()
-        
-        def _jax_wrapped_calc_ceil_approx(x, param):
-            return -jax_floor(-x, param) 
-        
-        return _jax_wrapped_calc_ceil_approx, jax_param
+    def sgn(self, id, init_params):
+        return self.comparison.sgn(id, init_params)
     
-    def mod(self):
-        jax_floor, jax_param = self.floor()
+    def floor(self, id, init_params):
+        return self.rounding.floor(id, init_params)
         
-        def _jax_wrapped_calc_mod_approx(x, y, param):
-            return x - y * jax_floor(x / y, param)
-        
-        return _jax_wrapped_calc_mod_approx, jax_param
+    def round(self, id, init_params):
+        return self.rounding.round(id, init_params)
     
-    def div(self):
-        jax_floor, jax_param = self.floor()
-        
-        def _jax_wrapped_calc_div_approx(x, y, param):
-            return jax_floor(x / y, param)
-        
-        return _jax_wrapped_calc_div_approx, jax_param
+    def ceil(self, id, init_params):
+        _floor = self.rounding.floor(id, init_params)
+        def _jax_wrapped_calc_ceil_approx(x, params):
+            neg_floor, params = _floor(-x, params) 
+            return -neg_floor, params        
+        return _jax_wrapped_calc_ceil_approx
     
-    def sqrt(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: sqrt(x) --> sqrt(x + eps)')
-        
-        def _jax_wrapped_calc_sqrt_approx(x, param):
-            return jnp.sqrt(x + self.eps)
-        
-        return _jax_wrapped_calc_sqrt_approx, None
+    def div(self, id, init_params):
+        _floor = self.rounding.floor(id, init_params)        
+        def _jax_wrapped_calc_div_approx(x, y, params):
+            return _floor(x / y, params)
+        return _jax_wrapped_calc_div_approx
+    
+    def mod(self, id, init_params):
+        _div = self.div(id, init_params)        
+        def _jax_wrapped_calc_mod_approx(x, y, params):
+            div, params = _div(x, y, params)
+            return x - y * div, params        
+        return _jax_wrapped_calc_mod_approx
+    
+    def sqrt(self, id, init_params):      
+        def _jax_wrapped_calc_sqrt_approx(x, params):
+            return jnp.sqrt(x + self.eps), params   
+        return _jax_wrapped_calc_sqrt_approx
     
     # ===========================================================================
     # indexing
     # ===========================================================================
      
-    @staticmethod
-    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)
-        literals = jnp.broadcast_to(literals, shape=shape)
-        return literals
-    
-    def argmax(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'argmax(x) --> sum(i * softmax(x[i]))')
-            
-        debias = 'argmax' in self.debias
-        
-        # https://arxiv.org/abs/2110.05651
-        def _jax_wrapped_calc_argmax_approx(x, axis, param):
-            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:
-                hard_sample = jnp.argmax(x, axis=axis)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'argmax')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_argmax_approx, new_param
+    def argmax(self, id, init_params): 
+        return self.comparison.argmax(id, init_params)
     
-    def argmin(self):
-        jax_argmax, jax_param = self.argmax()
-        
+    def argmin(self, id, init_params):
+        _argmax = self.argmax(id, init_params)
         def _jax_wrapped_calc_argmin_approx(x, axis, param):
-            return jax_argmax(-x, axis, param)
-        
-        return _jax_wrapped_calc_argmin_approx, jax_param
+            return _argmax(-x, axis, param)        
+        return _jax_wrapped_calc_argmin_approx
     
     # ===========================================================================
     # control flow
     # ===========================================================================
      
-    def control_if(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'if c then a else b --> c * a + (1 - c) * b')
-        
-        debias = 'if' in self.debias
-        
-        def _jax_wrapped_calc_if_approx(c, a, b, param):
-            sample = c * a + (1.0 - c) * b
-            if debias:
-                hard_sample = jnp.where(c > 0.5, a, b)
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        return _jax_wrapped_calc_if_approx, None
-    
-    def control_switch(self):
-        if self.verbose:
-            raise_warning('Using the replacement rule: '
-                          'switch(pred) { cases } --> '
-                          'sum(cases[i] * softmax(-(pred - i)^2))')   
-            
-        debias = 'switch' in self.debias
-        
-        def _jax_wrapped_calc_switch_approx(pred, cases, param):
-            literals = FuzzyLogic.enumerate_literals(cases.shape, axis=0)
-            pred = jnp.broadcast_to(pred[jnp.newaxis, ...], shape=cases.shape)
-            proximity = -jnp.square(pred - literals)
-            soft_case = jax.nn.softmax(param * proximity, axis=0)
-            sample = jnp.sum(cases * soft_case, axis=0)
-            if debias:
-                hard_case = jnp.argmax(proximity, axis=0)[jnp.newaxis, ...]      
-                hard_sample = jnp.take_along_axis(cases, hard_case, axis=0)[0, ...]
-                sample += jax.lax.stop_gradient(hard_sample - sample)
-            return sample
-        
-        tags = ('weight', 'switch')
-        new_param = (tags, self.weight)
-        return _jax_wrapped_calc_switch_approx, new_param
+    def control_if(self, id, init_params):
+        return self.control.if_then_else(id, init_params)
+    
+    def control_switch(self, id, init_params):
+        return self.control.switch(id, init_params)
     
     # ===========================================================================
     # random variables
     # ===========================================================================
      
-    def discrete(self):
-        return self.sampling.discrete(self)
+    def discrete(self, id, init_params):
+        return self.sampling.discrete(id, init_params, self)
     
-    def bernoulli(self):
-        return self.sampling.bernoulli(self)
+    def bernoulli(self, id, init_params):
+        return self.sampling.bernoulli(id, init_params, self)
     
-    def poisson(self):
-        return self.sampling.poisson(self)
+    def poisson(self, id, init_params):
+        return self.sampling.poisson(id, init_params, self)
     
-    def geometric(self):
-        return self.sampling.geometric(self)
+    def geometric(self, id, init_params):
+        return self.sampling.geometric(id, init_params, self)
 
 
 # ===========================================================================
@@ -733,105 +957,121 @@ class FuzzyLogic:
 #
 # ===========================================================================
 
-logic = FuzzyLogic()
-w = 1000.0
+logic = FuzzyLogic(comparison=SigmoidComparison(10000.0),
+                   rounding=SoftRounding(10000.0),
+                   control=SoftControlFlow(10000.0))
 
 
 def _test_logical():
     print('testing logical')
-    _and, _ = logic.logical_and()
-    _not, _ = logic.logical_not()
-    _gre, _ = logic.greater()
-    _or, _ = logic.logical_or()
-    _if, _ = logic.control_if()
-    
+    init_params = {}
+    _and = logic.logical_and(0, init_params)
+    _not = logic.logical_not(1, init_params)
+    _gre = logic.greater(2, init_params)
+    _or = logic.logical_or(3, init_params)
+    _if = logic.control_if(4, init_params)
+    print(init_params)
+
     # https://towardsdatascience.com/emulating-logical-gates-with-a-neural-network-75c229ec4cc9
-    def test_logic(x1, x2):
-        q1 = _and(_gre(x1, 0, w), _gre(x2, 0, w), w)
-        q2 = _and(_not(_gre(x1, 0, w), w), _not(_gre(x2, 0, w), w), w)
-        cond = _or(q1, q2, w)
-        pred = _if(cond, +1, -1, w)
+    def test_logic(x1, x2, w):
+        q1, w = _gre(x1, 0, w)
+        q2, w = _gre(x2, 0, w)
+        q3, w = _and(q1, q2, w)
+        q4, w = _not(q1, w)
+        q5, w = _not(q2, w)
+        q6, w = _and(q4, q5, w)        
+        cond, w = _or(q3, q6, w)
+        pred, w = _if(cond, +1, -1, w)
         return pred
     
     x1 = jnp.asarray([1, 1, -1, -1, 0.1, 15, -0.5]).astype(float)
     x2 = jnp.asarray([1, -1, 1, -1, 10, -30, 6]).astype(float)
-    print(test_logic(x1, x2))
+    print(test_logic(x1, x2, init_params))    
 
 
 def _test_indexing():
     print('testing indexing')
-    _argmax, _ = logic.argmax()
-    _argmin, _ = logic.argmin()
+    init_params = {}
+    _argmax = logic.argmax(0, init_params)
+    _argmin = logic.argmin(1, init_params)
+    print(init_params)
 
-    def argmaxmin(x):
-        amax = _argmax(x, 0, w)
-        amin = _argmin(x, 0, w)
+    def argmaxmin(x, w):
+        amax, w = _argmax(x, 0, w)
+        amin, w = _argmin(x, 0, w)
         return amax, amin
         
     values = jnp.asarray([2., 3., 5., 4.9, 4., 1., -1., -2.])
-    amax, amin = argmaxmin(values)
+    amax, amin = argmaxmin(values, init_params)
     print(amax)
     print(amin)
 
 
 def _test_control():
     print('testing control')
-    _switch, _ = logic.control_switch()
+    init_params = {}
+    _switch = logic.control_switch(0, init_params)
+    print(init_params)
     
     pred = jnp.asarray(jnp.linspace(0, 2, 10))
     case1 = jnp.asarray([-10.] * 10)
     case2 = jnp.asarray([1.5] * 10)
     case3 = jnp.asarray([10.] * 10)
     cases = jnp.asarray([case1, case2, case3])
-    print(_switch(pred, cases, w))
+    switch, _ = _switch(pred, cases, init_params)
+    print(switch)
 
 
 def _test_random():
     print('testing random')
     key = random.PRNGKey(42)
-    _bernoulli, _ = logic.bernoulli()
-    _discrete, _ = logic.discrete()
-    _geometric, _ = logic.geometric()
+    init_params = {}
+    _bernoulli = logic.bernoulli(0, init_params)
+    _discrete = logic.discrete(1, init_params)
+    _geometric = logic.geometric(2, init_params)
+    print(init_params)
     
-    def bern(n):
+    def bern(n, w):
         prob = jnp.asarray([0.3] * n)
-        sample = _bernoulli(key, prob, w)
+        sample, _ = _bernoulli(key, prob, w)
         return sample
     
-    samples = bern(50000)
+    samples = bern(50000, init_params)
     print(jnp.mean(samples))
     
-    def disc(n):
+    def disc(n, w):
         prob = jnp.asarray([0.1, 0.4, 0.5])
         prob = jnp.tile(prob, (n, 1))
-        sample = _discrete(key, prob, w)
+        sample, _ = _discrete(key, prob, w)
         return sample
         
-    samples = disc(50000)
+    samples = disc(50000, init_params)
     samples = jnp.round(samples)
     print([jnp.mean(samples == i) for i in range(3)])
     
-    def geom(n):
+    def geom(n, w):
         prob = jnp.asarray([0.3] * n)
-        sample = _geometric(key, prob, w)
+        sample, _ = _geometric(key, prob, w)
         return sample
     
-    samples = geom(50000)
+    samples = geom(50000, init_params)
     print(jnp.mean(samples))
     
 
 def _test_rounding():
     print('testing rounding')
-    _floor, _ = logic.floor()
-    _ceil, _ = logic.ceil()
-    _round, _ = logic.round()
-    _mod, _ = logic.mod()
+    init_params = {}
+    _floor = logic.floor(0, init_params)
+    _ceil = logic.ceil(1, init_params)
+    _round = logic.round(2, init_params)
+    _mod = logic.mod(3, init_params)
+    print(init_params)
     
     x = jnp.asarray([2.1, 0.6, 1.99, -2.01, -3.2, -0.1, -1.01, 23.01, -101.99, 200.01])
-    print(_floor(x, w))
-    print(_ceil(x, w))
-    print(_round(x, w))
-    print(_mod(x, 2.0, w))
+    print(_floor(x, init_params)[0])
+    print(_ceil(x, init_params)[0])
+    print(_round(x, init_params)[0])
+    print(_mod(x, 2.0, init_params)[0])
 
 
 if __name__ == '__main__':