pyRDDLGym-jax 2.3__py3-none-any.whl → 2.4__py3-none-any.whl

pyRDDLGym_jax/__init__.py

@@ -1 +1 @@
-__version__ = '2.3'
+__version__ = '2.4'

pyRDDLGym_jax/core/compiler.py

@@ -471,8 +471,7 @@ class JaxRDDLCompiler:
         return printed
     
     def model_parameter_info(self) -> Dict[str, Dict[str, Any]]:
-        '''Returns a dictionary of additional information about model
-        parameters.'''
+        '''Returns a dictionary of additional information about model parameters.'''
         result = {}
         for (id, value) in self.model_params.items():
             expr_id = int(str(id).split('_')[0])
@@ -799,7 +798,7 @@ class JaxRDDLCompiler:
         elif n == 2 or (n >= 2 and op in {'*', '+'}):
             jax_exprs = [self._jax(arg, init_params) for arg in args]
             result = jax_exprs[0]
-            for i, jax_rhs in enumerate(jax_exprs[1:]):
+            for (i, jax_rhs) in enumerate(jax_exprs[1:]):
                 jax_op = valid_ops[op](f'{expr.id}_{op}{i}', init_params)
                 result = self._jax_binary(result, jax_rhs, jax_op, at_least_int=True)
             return result

pyRDDLGym_jax/core/logic.py

@@ -29,6 +29,8 @@
 #
 # ***********************************************************************
 
+
+from abc import ABCMeta, abstractmethod
 import traceback
 from typing import Callable, Dict, Tuple, Union
 
@@ -64,30 +66,35 @@ def enumerate_literals(shape: Tuple[int, ...], axis: int, dtype: type=jnp.int32)
 #
 # ===========================================================================
 
-class Comparison:
+class Comparison(metaclass=ABCMeta):
     '''Base class for approximate comparison operations.'''
     
+    @abstractmethod
     def greater_equal(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def greater(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def equal(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def sgn(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def argmax(self, id, init_params):
-        raise NotImplementedError
+        pass
     
 
 class SigmoidComparison(Comparison):
     '''Comparison operations approximated using sigmoid functions.'''
     
     def __init__(self, weight: float=10.0) -> None:
-        self.weight = weight
+        self.weight = float(weight)
         
     # https://arxiv.org/abs/2110.05651
     def greater_equal(self, id, init_params):
@@ -139,21 +146,23 @@ class SigmoidComparison(Comparison):
 #
 # ===========================================================================
 
-class Rounding:
+class Rounding(metaclass=ABCMeta):
     '''Base class for approximate rounding operations.'''
     
+    @abstractmethod
     def floor(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def round(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class SoftRounding(Rounding):
     '''Rounding operations approximated using soft operations.'''
     
     def __init__(self, weight: float=10.0) -> None:
-        self.weight = weight
+        self.weight = float(weight)
         
     # https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/bijectors/Softfloor
     def floor(self, id, init_params):
@@ -189,11 +198,12 @@ class SoftRounding(Rounding):
 #
 # ===========================================================================
 
-class Complement:
+class Complement(metaclass=ABCMeta):
     '''Base class for approximate logical complement operations.'''
     
+    @abstractmethod
     def __call__(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class StandardComplement(Complement):
@@ -222,16 +232,18 @@ class StandardComplement(Complement):
 # https://www.sciencedirect.com/science/article/abs/pii/016501149190171L
 # ===========================================================================
 
-class TNorm:
+class TNorm(metaclass=ABCMeta):
     '''Base class for fuzzy differentiable t-norms.'''
     
+    @abstractmethod
     def norm(self, id, init_params):
         '''Elementwise t-norm of x and y.'''
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def norms(self, id, init_params):
         '''T-norm computed for tensor x along axis.'''
-        raise NotImplementedError
+        pass
     
 
 class ProductTNorm(TNorm):
@@ -339,26 +351,32 @@ class YagerTNorm(TNorm):
 #
 # ===========================================================================
 
-class RandomSampling:
+class RandomSampling(metaclass=ABCMeta):
     '''Describes how non-reparameterizable random variables are sampled.'''
     
+    @abstractmethod
     def discrete(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def poisson(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def binomial(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def negative_binomial(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def geometric(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def bernoulli(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
     def __str__(self) -> str:
         return 'RandomSampling'
@@ -603,21 +621,23 @@ class Determinization(RandomSampling):
 #
 # ===========================================================================
 
-class ControlFlow:
+class ControlFlow(metaclass=ABCMeta):
     '''A base class for control flow, including if and switch statements.'''
     
+    @abstractmethod
     def if_then_else(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def switch(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class SoftControlFlow(ControlFlow):
     '''Soft control flow using a probabilistic interpretation.'''
     
     def __init__(self, weight: float=10.0) -> None:
-        self.weight = weight
+        self.weight = float(weight)
         
     @staticmethod
     def _jax_wrapped_calc_if_then_else_soft(c, a, b, params):
@@ -651,15 +671,15 @@ class SoftControlFlow(ControlFlow):
 # ===========================================================================
 
 
-class Logic:
+class Logic(metaclass=ABCMeta):
     '''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 summarize_hyperparameters(self) -> str:
+        return (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.'''
@@ -765,119 +785,150 @@ class Logic:
     # ===========================================================================
     # logical operators
     # ===========================================================================
-     
+    
+    @abstractmethod
     def logical_and(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def logical_not(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def logical_or(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def xor(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def implies(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def equiv(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def forall(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def exists(self, id, init_params):    
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # comparison operators
     # ===========================================================================
     
+    @abstractmethod
     def greater_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def greater(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def less_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def less(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def not_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # special functions
     # ===========================================================================
      
+    @abstractmethod
     def sgn(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def floor(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def round(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def ceil(self, id, init_params):
-        raise NotImplementedError 
+        pass 
     
+    @abstractmethod
     def div(self, id, init_params):
-        raise NotImplementedError 
+        pass 
     
+    @abstractmethod
     def mod(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def sqrt(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # indexing
     # ===========================================================================
-     
+    
+    @abstractmethod
     def argmax(self, id, init_params):   
-        raise NotImplementedError  
+        pass
     
+    @abstractmethod
     def argmin(self, id, init_params):   
-        raise NotImplementedError     
+        pass
     
     # ===========================================================================
     # control flow
     # ===========================================================================
      
+    @abstractmethod
     def control_if(self, id, init_params):
-        raise NotImplementedError
+        pass
         
+    @abstractmethod
     def control_switch(self, id, init_params):
-        raise NotImplementedError
+        pass
     
     # ===========================================================================
     # random variables
     # ===========================================================================
      
+    @abstractmethod
     def discrete(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def bernoulli(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def poisson(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def geometric(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def binomial(self, id, init_params):
-        raise NotImplementedError
+        pass
 
+    @abstractmethod
     def negative_binomial(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class ExactLogic(Logic):
@@ -1109,8 +1160,8 @@ class FuzzyLogic(Logic):
                 f'    underflow_tol={self.eps}\n'
                 f'    use_64_bit   ={self.use64bit}\n')
 
-    def summarize_hyperparameters(self) -> None:
-        print(self.__str__())
+    def summarize_hyperparameters(self) -> str:
+        return self.__str__()
         
     # ===========================================================================
     # logical operators