Adding TorchScalarVariable and TorchNDVariable

lume_torch/base.py

@@ -10,7 +10,13 @@
 import numpy as np
 from pydantic import BaseModel, ConfigDict, field_validator
 
-from lume_torch.variables import ScalarVariable, get_variable, ConfigEnum
+from lume_torch.variables import (
+    TorchScalarVariable,
+    get_variable,
+    ConfigEnum,
+    DistributionVariable,
+    TorchNDVariable,
+)
 from lume_torch.utils import (
     try_import_module,
     verify_unique_variable_names,
@@ -34,6 +40,11 @@
     np.float64: lambda x: float(x),
 }
 
+# Add torch.Tensor encoder if torch is available
+torch = try_import_module("torch")
+if torch is not None:
+    JSON_ENCODERS[torch.Tensor] = lambda x: x.tolist()
+
 
 def process_torch_module(
     module,
@@ -341,9 +352,9 @@ class LUMETorch(BaseModel, ABC):
 
     Attributes
     ----------
-    input_variables : list of ScalarVariable
+    input_variables : list of TorchScalarVariable
         List defining the input variables and their order.
-    output_variables : list of ScalarVariable
+    output_variables : list of TorchScalarVariable
         List defining the output variables and their order.
     input_validation_config : dict of str to ConfigEnum, optional
         Determines the behavior during input validation by specifying the validation
@@ -378,8 +389,10 @@ class LUMETorch(BaseModel, ABC):
 
     """
 
-    input_variables: list[ScalarVariable]
-    output_variables: list[ScalarVariable]
+    input_variables: list[Union[TorchScalarVariable, TorchNDVariable]]
+    output_variables: list[
+        Union[TorchScalarVariable, TorchNDVariable, DistributionVariable]
+    ]
     input_validation_config: Optional[dict[str, ConfigEnum]] = None
     output_validation_config: Optional[dict[str, ConfigEnum]] = None
 
@@ -396,7 +409,7 @@ def validate_input_variables(cls, value):
 
         Returns
         -------
-        list of ScalarVariable
+        list of TorchScalarVariable
             List of validated variable instances.
 
         Raises
@@ -411,7 +424,14 @@ def validate_input_variables(cls, value):
                 if isinstance(val, dict):
                     variable_class = get_variable(val["variable_class"])
                     new_value.append(variable_class(name=name, **val))
-                elif isinstance(val, ScalarVariable):
+                elif isinstance(
+                    val,
+                    (
+                        TorchScalarVariable,
+                        TorchNDVariable,
+                        DistributionVariable,
+                    ),
+                ):
                     new_value.append(val)
                 else:
                     raise TypeError(f"type {type(val)} not supported")
@@ -510,13 +530,18 @@ def input_validation(self, input_dict: dict[str, Any]) -> dict[str, Any]:
 
         """
         for name, value in input_dict.items():
-            _config = (
-                "none"
-                if self.input_validation_config is None
-                else self.input_validation_config.get(name)
-            )
-            var = self.input_variables[self.input_names.index(name)]
-            var.validate_value(value, config=_config)
+            if name in self.input_names:
+                _config = (
+                    None
+                    if self.input_validation_config is None
+                    else self.input_validation_config.get(name)
+                )
+                var = self.input_variables[self.input_names.index(name)]
+                var.validate_value(value, config=_config)
+            else:
+                raise ValueError(
+                    f"Input variable {name} not found in model input variables."
+                )
         return input_dict
 
     def output_validation(self, output_dict: dict[str, Any]) -> dict[str, Any]:

lume_torch/models/ensemble.py

@@ -10,13 +10,13 @@
 from torch.distributions import Normal
 from torch.distributions.distribution import Distribution as TDistribution
 
-from lume_torch.models.prob_model_base import ProbModelBaseModel
+from lume_torch.models.prob_model_base import ProbabilisticBaseModel
 from lume_torch.models.torch_model import TorchModel
 
 logger = logging.getLogger(__name__)
 
 
-class NNEnsemble(ProbModelBaseModel):
+class NNEnsemble(ProbabilisticBaseModel):
     """LUME-model class for neural network ensembles.
 
     This class allows for the evaluation of multiple neural network models as an
@@ -43,9 +43,9 @@ def __init__(self, *args, **kwargs):
         Parameters
         ----------
         *args
-            Positional arguments forwarded to :class:`ProbModelBaseModel`.
+            Positional arguments forwarded to :class:`ProbabilisticBaseModel`.
         **kwargs
-            Keyword arguments forwarded to :class:`ProbModelBaseModel`.
+            Keyword arguments forwarded to :class:`ProbabilisticBaseModel`.
 
         Notes
         -----
@@ -111,7 +111,7 @@ def _get_predictions(
     ) -> dict[str, TDistribution]:
         """Get the predictions of the ensemble of models.
 
-        This implements the abstract method from :class:`ProbModelBaseModel` by
+        This implements the abstract method from :class:`ProbabilisticBaseModel` by
         evaluating each model in the ensemble and aggregating their outputs.
 
         Parameters