Models

LUME-torch provides several model classes for different use cases, from simple custom models to advanced probabilistic models.

Overview

All LUME-torch models inherit from LUMETorch and provide:

• Consistent input/output interface using dictionaries
• Variable validation
• Configuration file support (YAML)
• Serialization/deserialization

Base Model

LUMETorch

The foundation for all LUME models.

lume_torch.base.LUMETorch

Bases: BaseModel, ABC

Abstract base class for models using lume-torch variables.

Inheriting classes must define the _evaluate method and variable names must be unique. Models built using this framework will be compatible with the lume-epics EPICS server and associated tools.

Attributes

input_variables : list of TorchScalarVariable List defining the input variables and their order. 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 config for each input variable: {var_name: value}. Value can be "warn", "error", or "none". output_validation_config : dict of str to ConfigEnum, optional Determines the behavior during output validation by specifying the validation config for each output variable: {var_name: value}. Value can be "warn", "error", or "none".

Methods

evaluate(input_dict, kwargs) Main evaluation function that validates inputs, calls _evaluate, and validates outputs. input_validation(input_dict) Validates input dictionary values against input variable specifications. output_validation(output_dict) Validates output dictionary values against output variable specifications. yaml(base_key="", file_prefix="", save_models=False, save_jit=False) Serializes the model to a YAML formatted string. dump(file, base_key="", save_models=True, save_jit=False) Saves model configuration and associated files to disk. from_file(filename) Class method to load a model from a YAML file. from_yaml(yaml_obj) Class method to load a model from a YAML string or file object. register_to_mlflow(artifact_path, kwargs) Registers the model to MLflow for experiment tracking.

Notes

Subclasses must implement the abstract method _evaluate(input_dict, **kwargs) which performs the actual model computation.

Source code in lume_torch/base.py

class LUMETorch(BaseModel, ABC):
    """Abstract base class for models using lume-torch variables.

    Inheriting classes must define the _evaluate method and variable names must be unique.
    Models built using this framework will be compatible with the lume-epics EPICS server and associated tools.

    Attributes
    ----------
    input_variables : list of TorchScalarVariable
        List defining the input variables and their order.
    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
        config for each input variable: {var_name: value}. Value can be "warn", "error", or "none".
    output_validation_config : dict of str to ConfigEnum, optional
        Determines the behavior during output validation by specifying the validation
        config for each output variable: {var_name: value}. Value can be "warn", "error", or "none".

    Methods
    -------
    evaluate(input_dict, **kwargs)
        Main evaluation function that validates inputs, calls _evaluate, and validates outputs.
    input_validation(input_dict)
        Validates input dictionary values against input variable specifications.
    output_validation(output_dict)
        Validates output dictionary values against output variable specifications.
    yaml(base_key="", file_prefix="", save_models=False, save_jit=False)
        Serializes the model to a YAML formatted string.
    dump(file, base_key="", save_models=True, save_jit=False)
        Saves model configuration and associated files to disk.
    from_file(filename)
        Class method to load a model from a YAML file.
    from_yaml(yaml_obj)
        Class method to load a model from a YAML string or file object.
    register_to_mlflow(artifact_path, **kwargs)
        Registers the model to MLflow for experiment tracking.

    Notes
    -----
    Subclasses must implement the abstract method `_evaluate(input_dict, **kwargs)` which performs
    the actual model computation.

    """

    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

    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)

    @field_validator("input_variables", "output_variables", mode="before")
    def validate_input_variables(cls, value):
        """Validates and converts input/output variables to proper format.

        Parameters
        ----------
        value : dict or list
            Variables as dictionary or list to validate and convert.

        Returns
        -------
        list of TorchScalarVariable
            List of validated variable instances.

        Raises
        ------
        TypeError
            If variable type is not supported.

        """
        new_value = []
        if isinstance(value, dict):
            for name, val in value.items():
                if isinstance(val, dict):
                    variable_class = get_variable(val["variable_class"])
                    new_value.append(variable_class(name=name, **val))
                elif isinstance(
                    val,
                    (
                        TorchScalarVariable,
                        TorchNDVariable,
                        DistributionVariable,
                    ),
                ):
                    new_value.append(val)
                else:
                    raise TypeError(f"type {type(val)} not supported")
        elif isinstance(value, list):
            new_value = value
        return new_value

    def __init__(self, *args, **kwargs):
        """Initializes LUMETorch.

        Parameters
        ----------
        *args : dict, str, or os.PathLike
            Accepts a single argument which is the model configuration as dictionary, YAML or JSON
            formatted string or file path.
        **kwargs
            See class attributes.

        Raises
        ------
        ValueError
            If both YAML config and keyword arguments are provided, or if more than one
            positional argument is provided.

        """
        if len(args) == 1:
            if len(kwargs) > 0:
                logger.error("Cannot specify both YAML config and keyword arguments")
                raise ValueError(
                    "Cannot specify YAML string and keyword arguments for LUMETorch init."
                )
            logger.debug("Initializing model from configuration")
            super().__init__(**parse_config(args[0], type(self).model_fields))
        elif len(args) > 1:
            logger.error(f"Too many positional arguments: {len(args)}")
            raise ValueError(
                "Arguments to LUMETorch must be either a single YAML string "
                "or keyword arguments passed directly to pydantic."
            )
        else:
            logger.debug("Initializing model from keyword arguments")
            super().__init__(**kwargs)

        logger.info(
            f"Initialized {self.__class__.__name__} with {len(self.input_variables)} inputs and {len(self.output_variables)} outputs"
        )

    @field_validator("input_variables", "output_variables")
    def unique_variable_names(cls, value):
        verify_unique_variable_names(value)
        return value

    @property
    def input_names(self) -> list[str]:
        return [var.name for var in self.input_variables]

    @property
    def output_names(self) -> list[str]:
        return [var.name for var in self.output_variables]

    @property
    def default_input_validation_config(self) -> dict[str, ConfigEnum]:
        """Determines default behavior during input validation (if input_validation_config is None)."""
        return {var.name: var.default_validation_config for var in self.input_variables}

    @property
    def default_output_validation_config(self) -> dict[str, ConfigEnum]:
        """Determines default behavior during output validation (if output_validation_config is None)."""
        return {
            var.name: var.default_validation_config for var in self.output_variables
        }

    def evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
        """Main evaluation function, child classes must implement the _evaluate method."""
        self._validate_dict_keys(input_dict, dict_name="input")
        validated_input_dict = self.input_validation(input_dict)
        output_dict = self._evaluate(validated_input_dict, **kwargs)
        self._validate_dict_keys(output_dict, dict_name="output")
        self.output_validation(output_dict)
        return output_dict

    @abstractmethod
    def _evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
        pass

    def _validate_dict_keys(self, in_dict, dict_name="input"):
        """
        Validates that the keys in the input dictionary are a subset of the valid variable names.
        """
        valid_keys = self.input_names if dict_name == "input" else self.output_names
        valid_names = {name for name in valid_keys}
        invalid_keys = set(in_dict.keys()) - valid_names
        if invalid_keys:
            raise ValueError(
                f"Unknown {dict_name} variable(s): {sorted(invalid_keys)}. "
                f"Valid variables are: {sorted(valid_names)}"
            )

    def input_validation(self, input_dict: dict[str, Any]) -> dict[str, Any]:
        """Validates input dictionary values against input variable specifications.

        Parameters
        ----------
        input_dict : dict of str to Any
            Dictionary of input variable names to values.

        Returns
        -------
        dict of str to Any
            Validated input dictionary.

        """
        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)

        return input_dict

    def output_validation(self, output_dict: dict[str, Any]) -> dict[str, Any]:
        """Validates output dictionary values against output variable specifications.

        Parameters
        ----------
        output_dict : dict of str to Any
            Dictionary of output variable names to values.

        Returns
        -------
        dict of str to Any
            Validated output dictionary.

        Raises
        ------
        ValueError
            If ``output_dict`` contains a name not found in the model's output variables.

        """
        for name, value in output_dict.items():
            _config = (
                None
                if self.output_validation_config is None
                else self.output_validation_config.get(name)
            )
            var = self.output_variables[self.output_names.index(name)]
            var.validate_value(value, config=_config)
        return output_dict

    def to_json(self, **kwargs) -> str:
        """Serializes the model to a JSON formatted string.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments for serialization (base_key, file_prefix, save_models, save_jit).

        Returns
        -------
        str
            JSON formatted string defining the model.

        """
        return json_dumps(self, **kwargs)

    def model_dump(self, **kwargs) -> dict[str, Any]:
        """Dumps the model configuration as a dictionary.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments for Pydantic's model_dump.

        Returns
        -------
        dict of str to Any
            Dictionary containing the model configuration including model_class name.

        """
        config = super().model_dump(**kwargs)
        config["input_variables"] = [var.model_dump() for var in self.input_variables]
        config["output_variables"] = [var.model_dump() for var in self.output_variables]
        return {"model_class": self.__class__.__name__} | config

    def json(self, **kwargs) -> str:
        """Serializes the model to a JSON formatted string.

        Parameters
        ----------
        **kwargs
            Additional keyword arguments for serialization.

        Returns
        -------
        str
            JSON formatted string defining the model.

        """
        result = self.to_json(**kwargs)
        config = json.loads(result)
        return json.dumps(config)

    def yaml(
        self,
        base_key: str = "",
        file_prefix: str = "",
        save_models: bool = False,
        save_jit: bool = False,
    ) -> str:
        """Serializes the object and returns a YAML formatted string defining the model.

        Parameters
        ----------
        base_key : str, optional
            Base key for serialization.
        file_prefix : str, optional
            Prefix for generated filenames.
        save_models : bool, optional
            Determines whether models are saved to file.
        save_jit : bool, optional
            Determines whether the model is saved as TorchScript.

        Returns
        -------
        str
            YAML formatted string defining the model.

        """
        output = json.loads(
            self.to_json(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )
        s = yaml.dump(output, default_flow_style=None, sort_keys=False)
        return s

    def dump(
        self,
        file: Union[str, os.PathLike],
        base_key: str = "",
        save_models: bool = True,
        save_jit: bool = False,
    ):
        """Returns and optionally saves YAML formatted string defining the model.

        Parameters
        ----------
        file : str or os.PathLike
            File path to which the YAML formatted string and corresponding files are saved.
        base_key : str, optional
            Base key for serialization.
        save_models : bool, optional
            Determines whether models are saved to file.
        save_jit : bool, optional
            Determines whether the model is saved as TorchScript.

        """
        logger.info(f"Dumping model configuration to: {file}")
        if save_models:
            logger.debug("Saving model files alongside configuration")
        if save_jit:
            logger.debug("Saving models as TorchScript (JIT)")
        file_prefix = os.path.splitext(os.path.abspath(file))[0]
        with open(file, "w") as f:
            f.write(
                self.yaml(
                    base_key=base_key,
                    file_prefix=file_prefix,
                    save_models=save_models,
                    save_jit=save_jit,
                )
            )

    @classmethod
    def from_file(cls, filename: str):
        """Loads a model from a YAML file.

        Parameters
        ----------
        filename : str
            Path to the YAML file containing the model configuration.

        Returns
        -------
        LUMETorch
            Instance of the model loaded from the file.

        Raises
        ------
        OSError
            If the file does not exist.

        """
        if not os.path.exists(filename):
            raise OSError(f"File {filename} is not found.")
        with open(filename, "r") as file:
            return cls.from_yaml(file)

    @classmethod
    def from_yaml(cls, yaml_obj: str | TextIOWrapper):
        """Loads a model from a YAML string or file object.

        Parameters
        ----------
        yaml_obj : str or TextIOWrapper
            YAML formatted string or file object containing the model configuration.

        Returns
        -------
        LUMETorch
            Instance of the model loaded from the YAML configuration.

        """
        return cls.model_validate(parse_config(yaml_obj, cls.model_fields))

    def register_to_mlflow(
        self,
        artifact_path: str,
        registered_model_name: str | None = None,
        tags: dict[str, Any] | None = None,
        version_tags: dict[str, Any] | None = None,
        alias: str | None = None,
        run_name: str | None = None,
        log_model_dump: bool = True,
        save_jit: bool = False,
        **kwargs,
    ):
        """Registers the model to MLflow if mlflow is installed.

        Each time this function is called, a new version of the model is created. The model is saved to the
        tracking server or local directory, depending on the MLFLOW_TRACKING_URI.

        If no tracking server is set up, data and artifacts are saved directly under your current directory. To set up
        a tracking server, set the environment variable MLFLOW_TRACKING_URI, e.g. a local port/path. See
        https://mlflow.org/docs/latest/getting-started/intro-quickstart/ for more info.

        Parameters
        ----------
        artifact_path : str
            Path to store the model in MLflow.
        registered_model_name : str or None, optional
            Name of the registered model in MLflow.
        tags : dict of str to Any or None, optional
            Tags to add to the MLflow model.
        version_tags : dict of str to Any or None, optional
            Tags to add to this MLflow model version.
        alias : str or None, optional
            Alias to add to this MLflow model version.
        run_name : str or None, optional
            Name of the MLflow run.
        log_model_dump : bool, optional
            Whether to log the model dump files as artifacts.
        save_jit : bool, optional
            Whether to save the model as TorchScript when calling model.dump, if log_model_dump=True.
        **kwargs
            Additional arguments for mlflow.pyfunc.log_model.

        Returns
        -------
        mlflow.models.model.ModelInfo
            Model info metadata.

        """
        return register_model(
            self,
            artifact_path,
            registered_model_name,
            tags,
            version_tags,
            alias,
            run_name,
            log_model_dump,
            save_jit,
            **kwargs,
        )

__init__(*args, **kwargs)

Initializes LUMETorch.

Parameters

args : dict, str, or os.PathLike Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path. *kwargs See class attributes.

Raises

ValueError If both YAML config and keyword arguments are provided, or if more than one positional argument is provided.

Source code in lume_torch/base.py

def __init__(self, *args, **kwargs):
    """Initializes LUMETorch.

    Parameters
    ----------
    *args : dict, str, or os.PathLike
        Accepts a single argument which is the model configuration as dictionary, YAML or JSON
        formatted string or file path.
    **kwargs
        See class attributes.

    Raises
    ------
    ValueError
        If both YAML config and keyword arguments are provided, or if more than one
        positional argument is provided.

    """
    if len(args) == 1:
        if len(kwargs) > 0:
            logger.error("Cannot specify both YAML config and keyword arguments")
            raise ValueError(
                "Cannot specify YAML string and keyword arguments for LUMETorch init."
            )
        logger.debug("Initializing model from configuration")
        super().__init__(**parse_config(args[0], type(self).model_fields))
    elif len(args) > 1:
        logger.error(f"Too many positional arguments: {len(args)}")
        raise ValueError(
            "Arguments to LUMETorch must be either a single YAML string "
            "or keyword arguments passed directly to pydantic."
        )
    else:
        logger.debug("Initializing model from keyword arguments")
        super().__init__(**kwargs)

    logger.info(
        f"Initialized {self.__class__.__name__} with {len(self.input_variables)} inputs and {len(self.output_variables)} outputs"
    )

evaluate(input_dict, **kwargs)

Main evaluation function, child classes must implement the _evaluate method.

Source code in lume_torch/base.py

def evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
    """Main evaluation function, child classes must implement the _evaluate method."""
    self._validate_dict_keys(input_dict, dict_name="input")
    validated_input_dict = self.input_validation(input_dict)
    output_dict = self._evaluate(validated_input_dict, **kwargs)
    self._validate_dict_keys(output_dict, dict_name="output")
    self.output_validation(output_dict)
    return output_dict

_evaluate(input_dict, **kwargs) abstractmethod

Source code in lume_torch/base.py

@abstractmethod
def _evaluate(self, input_dict: dict[str, Any], **kwargs) -> dict[str, Any]:
    pass

input_validation(input_dict)

Validates input dictionary values against input variable specifications.

Parameters

input_dict : dict of str to Any Dictionary of input variable names to values.

Returns

dict of str to Any Validated input dictionary.

Source code in lume_torch/base.py

def input_validation(self, input_dict: dict[str, Any]) -> dict[str, Any]:
    """Validates input dictionary values against input variable specifications.

    Parameters
    ----------
    input_dict : dict of str to Any
        Dictionary of input variable names to values.

    Returns
    -------
    dict of str to Any
        Validated input dictionary.

    """
    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)

    return input_dict

output_validation(output_dict)

Validates output dictionary values against output variable specifications.

Parameters

output_dict : dict of str to Any Dictionary of output variable names to values.

Returns

dict of str to Any Validated output dictionary.

Raises

ValueError If output_dict contains a name not found in the model's output variables.

Source code in lume_torch/base.py

def output_validation(self, output_dict: dict[str, Any]) -> dict[str, Any]:
    """Validates output dictionary values against output variable specifications.

    Parameters
    ----------
    output_dict : dict of str to Any
        Dictionary of output variable names to values.

    Returns
    -------
    dict of str to Any
        Validated output dictionary.

    Raises
    ------
    ValueError
        If ``output_dict`` contains a name not found in the model's output variables.

    """
    for name, value in output_dict.items():
        _config = (
            None
            if self.output_validation_config is None
            else self.output_validation_config.get(name)
        )
        var = self.output_variables[self.output_names.index(name)]
        var.validate_value(value, config=_config)
    return output_dict

dump(file, base_key='', save_models=True, save_jit=False)

Returns and optionally saves YAML formatted string defining the model.

Parameters

file : str or os.PathLike File path to which the YAML formatted string and corresponding files are saved. base_key : str, optional Base key for serialization. save_models : bool, optional Determines whether models are saved to file. save_jit : bool, optional Determines whether the model is saved as TorchScript.

Source code in lume_torch/base.py

def dump(
    self,
    file: Union[str, os.PathLike],
    base_key: str = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Returns and optionally saves YAML formatted string defining the model.

    Parameters
    ----------
    file : str or os.PathLike
        File path to which the YAML formatted string and corresponding files are saved.
    base_key : str, optional
        Base key for serialization.
    save_models : bool, optional
        Determines whether models are saved to file.
    save_jit : bool, optional
        Determines whether the model is saved as TorchScript.

    """
    logger.info(f"Dumping model configuration to: {file}")
    if save_models:
        logger.debug("Saving model files alongside configuration")
    if save_jit:
        logger.debug("Saving models as TorchScript (JIT)")
    file_prefix = os.path.splitext(os.path.abspath(file))[0]
    with open(file, "w") as f:
        f.write(
            self.yaml(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )

Creating Custom Models

To create a custom model, inherit from LUMETorch and implement _evaluate:

from lume_torch.base import LUMETorch
from lume_torch.variables import ScalarVariable


class MyModel(LUMETorch):
    """Custom model implementing specific logic."""

    def _evaluate(self, input_dict):
        """Implement your model logic here.

        Parameters
        ----------
        input_dict : dict
            Dictionary mapping input variable names to values

        Returns
        -------
        dict
            Dictionary mapping output variable names to values
        """
        x = input_dict["x"]
        y = input_dict["y"]
        return {
            "sum": x + y,
            "product": x * y
        }


# Create model with variables
model = MyModel(
    input_variables=[
        ScalarVariable(name="x", value_range=[0, 10]),
        ScalarVariable(name="y", value_range=[0, 10]),
    ],
    output_variables=[
        ScalarVariable(name="sum"),
        ScalarVariable(name="product"),
    ]
)

# Use the model
result = model({"x": 3.0, "y": 4.0})

PyTorch Models

TorchModel

Wrapper for PyTorch neural networks.

lume_torch.models.torch_model.TorchModel

Bases: LUMETorch

LUME-model class for torch models.

By default, the models are assumed to be fixed, so all gradient computation is deactivated and the model and transformers are put in evaluation mode.

Attributes

model : torch.nn.Module The underlying torch model. input_variables : list of TorchScalarVariable or TorchNDVariable List defining the input variables and their order. Supports both scalar variables and multi-dimensional array variables. output_variables : list of TorchScalarVariable or TorchNDVariable List defining the output variables and their order. input_transformers : list of ReversibleInputTransform, torch.nn.Linear, or Callable Transformer objects applied to the inputs before passing to the model. output_transformers : list of ReversibleInputTransform, torch.nn.Linear, or Callable Transformer objects applied to the outputs of the model. output_format : {"tensor", "variable", "raw"} Determines format of outputs. "tensor" returns tensors, "variable" and "raw" return scalars where possible. device : torch.device or str Device on which the model will be evaluated. Defaults to "cpu". fixed_model : bool If True, the model and transformers are put in evaluation mode and all gradient computation is deactivated. precision : {"double", "single"} Precision of the model, either "double" or "single".

Methods

evaluate(input_dict, **kwargs) Evaluate the model on a dictionary of inputs and return outputs. input_validation(input_dict) Validate and normalize the input dictionary before evaluation. output_validation(output_dict) Validate the output dictionary after evaluation. random_input(n_samples=1) Generate random inputs consistent with the input variable ranges. random_evaluate(n_samples=1) Evaluate the model on random inputs. to(device) Move the model, transformers, and default values to a given device.

Notes

When using TorchNDVariable inputs, all inputs must be TorchNDVariable. Mixing TorchScalarVariable and TorchNDVariable is not currently supported.

Source code in lume_torch/models/torch_model.py

class TorchModel(LUMETorch):
    """LUME-model class for torch models.

    By default, the models are assumed to be fixed, so all gradient computation
    is deactivated and the model and transformers are put in evaluation mode.

    Attributes
    ----------
    model : torch.nn.Module
        The underlying torch model.
    input_variables : list of TorchScalarVariable or TorchNDVariable
        List defining the input variables and their order. Supports both scalar
        variables and multi-dimensional array variables.
    output_variables : list of TorchScalarVariable or TorchNDVariable
        List defining the output variables and their order.
    input_transformers : list of ReversibleInputTransform, torch.nn.Linear, or Callable
        Transformer objects applied to the inputs before passing to the model.
    output_transformers : list of ReversibleInputTransform, torch.nn.Linear, or Callable
        Transformer objects applied to the outputs of the model.
    output_format : {"tensor", "variable", "raw"}
        Determines format of outputs. "tensor" returns tensors, "variable" and
        "raw" return scalars where possible.
    device : torch.device or str
        Device on which the model will be evaluated. Defaults to ``"cpu"``.
    fixed_model : bool
        If ``True``, the model and transformers are put in evaluation mode and
        all gradient computation is deactivated.
    precision : {"double", "single"}
        Precision of the model, either ``"double"`` or ``"single"``.

    Methods
    -------
    evaluate(input_dict, **kwargs)
        Evaluate the model on a dictionary of inputs and return outputs.
    input_validation(input_dict)
        Validate and normalize the input dictionary before evaluation.
    output_validation(output_dict)
        Validate the output dictionary after evaluation.
    random_input(n_samples=1)
        Generate random inputs consistent with the input variable ranges.
    random_evaluate(n_samples=1)
        Evaluate the model on random inputs.
    to(device)
        Move the model, transformers, and default values to a given device.

    Notes
    -----
    When using TorchNDVariable inputs, all inputs must be TorchNDVariable.
    Mixing TorchScalarVariable and TorchNDVariable is not currently supported.

    """

    model: torch.nn.Module
    input_transformers: list[
        Union[ReversibleInputTransform, torch.nn.Linear, Callable]
    ] = Field(default_factory=list)
    output_transformers: list[
        Union[ReversibleInputTransform, torch.nn.Linear, Callable]
    ] = Field(default_factory=list)
    output_format: str = "tensor"
    device: Union[torch.device, str] = "cpu"
    fixed_model: bool = True
    precision: str = "double"

    def __init__(self, *args, **kwargs):
        """Initializes TorchModel.

        Parameters
        ----------
        *args : dict, str, or Path
            Accepts a single argument which is the model configuration as dictionary, YAML or JSON
            formatted string or file path.
        **kwargs
            See class attributes.

        """
        super().__init__(*args, **kwargs)

        # dtype property sets precision across model and transformers
        self.dtype

        # fixed model: set full model in eval mode and deactivate all gradients
        if self.fixed_model:
            is_scripted = isinstance(self.model, torch.jit.ScriptModule)
            self.model.eval().requires_grad_(False) if not is_scripted else None
            for t in self.input_transformers + self.output_transformers:
                if isinstance(t, torch.nn.Module):
                    t.eval().requires_grad_(False)

        # ensure consistent device
        self.to(self.device)

    @property
    def dtype(self):
        if self.precision == "double":
            self._dtype = torch.double
        elif self.precision == "single":
            self._dtype = torch.float
        else:
            raise ValueError(
                f"Unknown precision {self.precision}, "
                f"expected one of ['double', 'single']."
            )
        self._set_precision(self._dtype)
        return self._dtype

    @property
    def _tkwargs(self):
        return {"device": self.device, "dtype": self.dtype}

    @field_validator("model", mode="before")
    @classmethod
    def validate_torch_model(
        cls, v: Union[str, os.PathLike, torch.nn.Module]
    ) -> torch.nn.Module:
        """Validate and load the torch model from file if needed.

        Parameters
        ----------
        v : str, os.PathLike, or torch.nn.Module
            Model or path to model file.

        Returns
        -------
        torch.nn.Module
            Loaded or validated torch model.

        Raises
        ------
        OSError
            If the model file does not exist.
        """
        if isinstance(v, (str, os.PathLike)):
            if os.path.exists(v):
                fname = v
                try:
                    v = torch.jit.load(v)
                    logger.info(f"Loaded TorchScript (JIT) model from file: {fname}")
                except RuntimeError:
                    v = torch.load(v, weights_only=False)
                    logger.info(f"Loaded PyTorch model from file: {fname}")
            else:
                logger.error(f"File {v} not found")
                raise OSError(f"File {v} is not found.")
        return v

    @field_validator("input_variables")
    @classmethod
    def verify_input_default_value(
        cls, value: list[Union[TorchScalarVariable, TorchNDVariable]]
    ) -> list[Union[TorchScalarVariable, TorchNDVariable]]:
        """Verify that input variables have the required default values.

        Parameters
        ----------
        value : list of TorchScalarVariable or TorchNDVariable
            Input variables to validate.

        Returns
        -------
        list of TorchScalarVariable or TorchNDVariable
            Validated input variables.

        Raises
        ------
        ValueError
            If any input variable is missing a default value.
        """
        for var in value:
            if var.default_value is None:
                logger.error(
                    f"Input variable {var.name} is missing required default value"
                )
                raise ValueError(
                    f"Input variable {var.name} must have a default value."
                )
        return value

    @field_validator("input_transformers", "output_transformers", mode="before")
    @classmethod
    def validate_transformers(cls, v: Union[list, str, os.PathLike]) -> list:
        """Validate and load transformers from files if needed.

        Parameters
        ----------
        v : list, str, or os.PathLike
            List of transformers or paths to transformer files.

        Returns
        -------
        list
            List of loaded transformers.

        Raises
        ------
        ValueError
            If transformers are not provided as a list.
        OSError
            If a transformer file does not exist.
        """
        if v is None:
            return []
        if not isinstance(v, list):
            logger.error(f"Transformers must be a list, got {type(v)}")
            raise ValueError("Transformers must be passed as list.")
        loaded_transformers = []
        for t in v:
            if isinstance(t, (str, os.PathLike)):
                if os.path.exists(t):
                    t = torch.load(t, weights_only=False)
                    logger.debug(f"Loaded transformer from file: {t}")
                else:
                    logger.error(f"Transformer file {t} not found")
                    raise OSError(f"File {t} is not found.")
            loaded_transformers.append(t)
        v = loaded_transformers
        return v

    @field_validator("output_format")
    @classmethod
    def validate_output_format(cls, v: str) -> str:
        """Validate the output format.

        Parameters
        ----------
        v : str
            Output format to validate.

        Returns
        -------
        str
            Validated output format.

        Raises
        ------
        ValueError
            If output format is not one of the supported formats.
        """
        supported_formats = ["tensor", "variable", "raw"]
        if v not in supported_formats:
            logger.error(
                f"Invalid output format {v}, expected one of {supported_formats}"
            )
            raise ValueError(
                f"Unknown output format {v}, expected one of {supported_formats}."
            )
        return v

    def _set_precision(self, value: torch.dtype):
        """Sets the precision of the model and transformers.

        Parameters
        ----------
        value : torch.dtype
            Dtype to set for the model and transformers.
        """
        self.model.to(dtype=value)
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.to(dtype=value)

    def _default_to_tensor(
        self, default_value: Union[torch.Tensor, float]
    ) -> torch.Tensor:
        """Convert a default value to a tensor with proper dtype and device.

        Parameters
        ----------
        default_value : torch.Tensor or float
            Default value to convert.

        Returns
        -------
        torch.Tensor
            Default value as a tensor with proper dtype and device.
        """
        if isinstance(default_value, torch.Tensor):
            return default_value.detach().clone().to(**self._tkwargs)
        else:
            return torch.tensor(default_value, **self._tkwargs)

    def _evaluate(
        self,
        input_dict: dict[str, Union[float, torch.Tensor]],
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Evaluate the model on the given input dictionary.

        Parameters
        ----------
        input_dict : dict of str to float or torch.Tensor
            Input dictionary on which to evaluate the model.

        Returns
        -------
        dict of str to float or torch.Tensor
            Dictionary of output variable names to values.

        """
        formatted_inputs = format_inputs(input_dict)
        input_tensor = self._arrange_inputs(formatted_inputs)
        input_tensor = self._transform_inputs(input_tensor)
        output_tensor = self.model(input_tensor)
        output_tensor = self._transform_outputs(output_tensor)
        parsed_outputs = self._parse_outputs(output_tensor)
        output_dict = self._prepare_outputs(parsed_outputs)
        return output_dict

    def input_validation(self, input_dict: dict[str, Union[float, torch.Tensor]]):
        """Validate the input dictionary before evaluation.

        Parameters
        ----------
        input_dict : dict of str to float or torch.Tensor
            Input dictionary to validate.

        Returns
        -------
        dict of str to float or torch.Tensor
            Validated input dictionary.

        """
        # type/dtype check on raw user-provided values (before tensor conversion)
        for var in self.input_variables:
            config = (
                None
                if self.input_validation_config is None
                or var.name not in self.input_validation_config
                else self.input_validation_config[var.name]
            )
            if var.name in input_dict:
                if var.read_only:
                    var.validate_value(var.default_value, config=config)
                    var.validate_read_only(input_dict[var.name], config=config)
                else:
                    var.validate_value(input_dict[var.name], config=config)
            else:
                # check all other default values in case of dynamic changes to defaults
                var.validate_value(var.default_value, config=config)

        # format inputs as tensors w/o changing the dtype
        formatted_inputs = format_inputs(input_dict)

        # cast tensors to expected dtype and device
        formatted_inputs = {
            k: v.to(**self._tkwargs) for k, v in formatted_inputs.items()
        }

        return formatted_inputs

    def output_validation(self, output_dict: dict[str, Union[float, torch.Tensor]]):
        """Validate the output dictionary after evaluation.

        Parameters
        ----------
        output_dict : dict of str to float or torch.Tensor
            Output dictionary to validate.

        """
        for var in self.output_variables:
            config = (
                None
                if self.output_validation_config is None
                or var.name not in self.output_validation_config
                or self.output_validation_config[var.name] is None
                else self.output_validation_config[var.name]
            )
            if var.name in output_dict:
                var.validate_value(output_dict[var.name], config=config)

    def random_input(self, n_samples: int = 1) -> dict[str, torch.Tensor]:
        """Generates random input(s) for the model.

        Parameters
        ----------
        n_samples : int, optional
            Number of random samples to generate.

        Returns
        -------
        dict of str to torch.Tensor
            Dictionary of input variable names to tensors.

        Notes
        -----
        For TorchScalarVariable inputs, generates random values within the variable's
        value_range. For TorchNDVariable inputs, repeats the default value for
        the requested number of samples.

        """
        input_dict = {}
        for var in self.input_variables:
            if isinstance(var, TorchScalarVariable):
                input_dict[var.name] = var.value_range[0] + torch.rand(
                    size=(n_samples,), **self._tkwargs
                ) * (var.value_range[1] - var.value_range[0])
            elif isinstance(var, TorchNDVariable):
                # For ND variables, repeat the default value for n_samples
                # Works for any dimensionality: 1D arrays, 2D matrices, 3D images, etc.
                default = self._default_to_tensor(var.default_value)
                # Add batch dim and repeat n_samples times (keeping original shape)
                # e.g., (3, 64, 64) -> (1, 3, 64, 64) -> (n_samples, 3, 64, 64)
                input_dict[var.name] = default.unsqueeze(0).repeat(
                    n_samples, *([1] * default.ndim)
                )
        return input_dict

    def random_evaluate(
        self, n_samples: int = 1
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Return random evaluations of the model.

        Parameters
        ----------
        n_samples : int, optional
            Number of random samples to evaluate.

        Returns
        -------
        dict of str to float or torch.Tensor
            Dictionary of variable names to outputs.

        """
        random_input = self.random_input(n_samples)
        return self.evaluate(random_input)

    def to(self, device: Union[torch.device, str]):
        """Update the device for the model, transformers and default values.

        Parameters
        ----------
        device : torch.device or str
            Device on which the model will be evaluated.

        """
        self.model.to(device)
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.to(device)
        self.device = device

    def insert_input_transformer(
        self, new_transformer: ReversibleInputTransform, loc: int
    ):
        """Insert an additional input transformer at the given location.

        Parameters
        ----------
        new_transformer : ReversibleInputTransform
            New transformer to add.
        loc : int
            Location where the new transformer shall be added to the
            transformer list.

        """
        self.input_transformers = (
            self.input_transformers[:loc]
            + [new_transformer]
            + self.input_transformers[loc:]
        )

    def insert_output_transformer(
        self, new_transformer: ReversibleInputTransform, loc: int
    ):
        """Inserts an additional output transformer at the given location.

        Parameters
        ----------
        new_transformer : ReversibleInputTransform
            New transformer to add.
        loc : int
            Location where the new transformer shall be added to the transformer list.

        """
        self.output_transformers = (
            self.output_transformers[:loc]
            + [new_transformer]
            + self.output_transformers[loc:]
        )

    def update_input_variables_to_transformer(
        self, transformer_loc: int
    ) -> list[TorchScalarVariable]:
        """Return input variables updated to the transformer at the given location.

        Updated are the value ranges and defaults of the input variables. This
        allows, for example, adding a calibration transformer and updating the
        input variable specification accordingly.

        Parameters
        ----------
        transformer_loc : int
            Index of the input transformer to adjust for.

        Returns
        -------
        list of TorchScalarVariable
            The updated input variables.

        """
        x_old = {
            "min": torch.tensor(
                [var.value_range[0] for var in self.input_variables], dtype=self.dtype
            ),
            "max": torch.tensor(
                [var.value_range[1] for var in self.input_variables], dtype=self.dtype
            ),
            "default": torch.tensor(
                [var.default_value for var in self.input_variables], dtype=self.dtype
            ),
        }
        x_new = {}
        for key, x in x_old.items():
            # Make at least 2D
            if x.ndim == 0:
                x = x.unsqueeze(0)
            if x.ndim == 1:
                x = x.unsqueeze(0)

            # compute previous limits at transformer location
            for i in range(transformer_loc):
                if isinstance(self.input_transformers[i], ReversibleInputTransform):
                    x = self.input_transformers[i].transform(x)
                else:
                    x = self.input_transformers[i](x)
            # untransform of transformer to adjust for
            if isinstance(
                self.input_transformers[transformer_loc], ReversibleInputTransform
            ):
                x = self.input_transformers[transformer_loc].untransform(x)
            elif isinstance(self.input_transformers[transformer_loc], torch.nn.Linear):
                w = self.input_transformers[transformer_loc].weight
                b = self.input_transformers[transformer_loc].bias
                x = torch.matmul((x - b), torch.linalg.inv(w.T))
            else:
                raise NotImplementedError(
                    f"Reverse transformation for type {type(self.input_transformers[transformer_loc])} is not supported."
                )
            # backtrack through transformers
            for transformer in self.input_transformers[:transformer_loc][::-1]:
                if isinstance(transformer, ReversibleInputTransform):
                    x = transformer.untransform(x)
                elif isinstance(transformer, torch.nn.Linear):
                    w, b = transformer.weight, transformer.bias
                    x = torch.matmul((x - b), torch.linalg.inv(w.T))
                else:
                    raise NotImplementedError(
                        f"Reverse transformation for type {type(transformer)} is not supported."
                    )

            x_new[key] = x
        updated_variables = deepcopy(self.input_variables)
        for i, var in enumerate(updated_variables):
            var.value_range = [x_new["min"][0][i].item(), x_new["max"][0][i].item()]
            var.default_value = x_new["default"][0][i].item()
        return updated_variables

    def _fill_default_inputs(
        self, input_dict: dict[str, torch.Tensor]
    ) -> dict[str, torch.Tensor]:
        """Fill missing input variables with default values.

        Parameters
        ----------
        input_dict : dict of str to torch.Tensor
            Dictionary of input variable names to tensors.

        Returns
        -------
        dict of str to torch.Tensor
            Dictionary of input variable names to tensors with default values
            for missing inputs.

        """
        for var in self.input_variables:
            if var.name not in input_dict.keys():
                input_dict[var.name] = self._default_to_tensor(var.default_value)

        return input_dict

    def _arrange_inputs(
        self, formatted_inputs: dict[str, torch.Tensor]
    ) -> torch.Tensor:
        """Enforces ordering, batching, and default filling of inputs.

        * If all inputs are `TorchNDVariable`, stacks them into shape `(batch, num_arrays, *array_shape)`.
        * If all inputs are `TorchScalarVariable`, concatenates them so the last dimension matches the number of inputs,
         broadcasting defaults as needed.
        * If a mix of array and scalar inputs is provided, raises `NotImplementedError`.
        * Missing inputs are filled with their default values before arranging.

        Parameters
        ----------
        formatted_inputs : dict of str to torch.Tensor
            Dictionary of input variable names to tensors.

        Returns
        -------
        torch.Tensor
            Ordered input tensor to be passed to the transformers.

        """
        contains_array = any(
            isinstance(v, TorchNDVariable) for v in self.input_variables
        )
        contains_scalar = any(
            isinstance(v, TorchScalarVariable) for v in self.input_variables
        )

        if contains_array and contains_scalar:
            raise NotImplementedError(
                "Mixing TorchScalarVariable and TorchNDVariable inputs is not supported."
            )

        # All TorchNDVariable
        if contains_array:
            tensor_list = []
            batch_shape = None
            for var in self.input_variables:
                if var.name in formatted_inputs:
                    value = formatted_inputs[var.name]
                else:
                    value = self._default_to_tensor(var.default_value)

                expected_sample_shape = tuple(var.shape)
                sample_ndim = len(expected_sample_shape)
                if value.shape[-sample_ndim:] != expected_sample_shape:
                    raise ValueError(
                        f"Input {var.name} has shape {value.shape}, "
                        f"expected sample shape {expected_sample_shape}"
                    )

                current_batch = value.shape[:-sample_ndim]
                if current_batch == torch.Size():
                    # No batch dim provided -> add singleton batch
                    value = value.unsqueeze(0)
                    current_batch = torch.Size([1])

                if batch_shape is None:
                    batch_shape = current_batch
                elif current_batch != batch_shape:
                    raise ValueError(
                        f"Inputs have inconsistent batch shapes: "
                        f"{batch_shape} vs {current_batch}"
                    )

                tensor_list.append(value.to(**self._tkwargs))

            stacked = torch.stack(tensor_list, dim=1)  # (batch, num_arrays, ...)
            logger.debug(f"Arranged ND inputs into tensor shape: {stacked.shape}")
            return stacked

        # All TorchScalarVariables
        default_list = []
        for var in self.input_variables:
            default_list.append(self._default_to_tensor(var.default_value))

        default_tensor = torch.cat([d.flatten() for d in default_list]).to(
            **self._tkwargs
        )

        if formatted_inputs:
            batch_shape = None
            for var in self.input_variables:
                if var.name in formatted_inputs:
                    value = formatted_inputs[var.name]
                    if value.ndim > 0 and value.shape[-1] in (0, 1):
                        batch_shape = value.shape[:-1]
                    else:
                        batch_shape = value.shape
                    break

            if batch_shape and len(batch_shape) > 0:
                expanded_shape = (*batch_shape, default_tensor.shape[0])
                input_tensor = (
                    default_tensor.unsqueeze(0).expand(expanded_shape).clone()
                )
            else:
                input_tensor = default_tensor.unsqueeze(0)

            current_idx = 0
            for var in self.input_variables:
                if var.name in formatted_inputs:
                    value = formatted_inputs[var.name]
                    if value.ndim > 0 and value.shape[-1] == 1:
                        input_tensor[..., current_idx] = value.squeeze(-1)
                    else:
                        input_tensor[..., current_idx] = value
                current_idx += 1
        else:
            input_tensor = default_tensor.unsqueeze(0)

        expected_features = len(self.input_variables)
        if input_tensor.shape[-1] != expected_features:
            raise ValueError(
                "Last dimension of input tensor doesn't match the expected number of features\n"
                f"received: {input_tensor.shape}, expected {expected_features} as the last dimension"
            )

        return input_tensor

    def _transform_inputs(self, input_tensor: torch.Tensor) -> torch.Tensor:
        """Applies transformations to the inputs.

        Parameters
        ----------
        input_tensor : torch.Tensor
            Ordered input tensor to be passed to the transformers.

        Returns
        -------
        torch.Tensor
            Tensor of transformed inputs to be passed to the model.

        """
        # Make at least 2D
        if input_tensor.ndim == 0:
            input_tensor = input_tensor.unsqueeze(0)
        if input_tensor.ndim == 1:
            input_tensor = input_tensor.unsqueeze(0)

        for transformer in self.input_transformers:
            if isinstance(transformer, ReversibleInputTransform):
                input_tensor = transformer.transform(input_tensor)
            else:
                input_tensor = transformer(input_tensor)
        return input_tensor

    def _transform_outputs(self, output_tensor: torch.Tensor) -> torch.Tensor:
        """(Un-)transform the model output tensor.

        Parameters
        ----------
        output_tensor : torch.Tensor
            Output tensor from the model.

        Returns
        -------
        torch.Tensor
            (Un-)transformed output tensor.

        """
        for transformer in self.output_transformers:
            if isinstance(transformer, ReversibleInputTransform):
                output_tensor = transformer.untransform(output_tensor)
            elif isinstance(transformer, torch.nn.Linear):
                w, b = transformer.weight, transformer.bias
                output_tensor = torch.matmul((output_tensor - b), torch.linalg.inv(w.T))
            else:
                # we assume anything else is provided as a callable
                output_tensor = transformer(output_tensor)
        return output_tensor

    def _parse_outputs(self, output_tensor: torch.Tensor) -> dict[str, torch.Tensor]:
        """Construct a dictionary from the model output tensor.

        Parameters
        ----------
        output_tensor : torch.Tensor
            (Un-)transformed output tensor from the model.

        Returns
        -------
        dict of str to torch.Tensor
            Dictionary of output variable names to (un-)transformed tensors.

        """
        parsed_outputs = {}

        # Check if all outputs are scalar variables
        all_scalars = all(
            isinstance(v, TorchScalarVariable) for v in self.output_variables
        )

        # Handle 0D and 1D tensors
        if output_tensor.dim() == 0:
            # 0D tensor - always add batch dimension at start
            output_tensor = output_tensor.unsqueeze(0)
        elif output_tensor.dim() == 1:
            # 1D tensor - interpretation depends on variable types
            if all_scalars and len(self.output_names) == 1:
                # For single scalar output, 1D means (batch,) -> should become (batch, 1)
                output_tensor = output_tensor.unsqueeze(-1)
            elif all_scalars and len(self.output_names) > 1:
                # For multiple scalar outputs, 1D means (features,) -> should become (1, features)
                output_tensor = output_tensor.unsqueeze(0)
            else:
                # For non-scalar outputs, default to adding batch dimension at start
                output_tensor = output_tensor.unsqueeze(0)

        if len(self.output_names) == 1:
            output = output_tensor
            # For scalar variables, ensure shape is (batch, 1) for single-sample batched outputs
            # or (batch, samples) for multi-sample outputs
            if all_scalars:
                if output.dim() == 2:
                    # Already 2D: could be (batch, 1) or (batch, samples) - keep as is
                    parsed_outputs[self.output_names[0]] = output
                elif output.dim() == 1:
                    # Shape is (batch,), reshape to (batch, 1)
                    parsed_outputs[self.output_names[0]] = output.unsqueeze(-1)
                else:
                    # 3D or higher dimensional - squeeze last dim if it's 1
                    # This handles multi-sample cases: (batch, samples, 1) -> (batch, samples)
                    if output.shape[-1] != 1:
                        parsed_outputs[self.output_names[0]] = output.unsqueeze(-1)
                    else:
                        # Shouldn't happen, but handle by squeezing all and adding feature dim
                        parsed_outputs[self.output_names[0]] = (
                            output.squeeze().unsqueeze(-1)
                            if output.squeeze().dim() > 0
                            else output.squeeze().unsqueeze(0).unsqueeze(-1)
                        )
            else:
                # For non-scalar outputs (NDVariable), keep original behavior
                parsed_outputs[self.output_names[0]] = output.squeeze()
        else:
            for idx, output_name in enumerate(self.output_names):
                output = output_tensor[..., idx]
                var = self.output_variables[idx]

                # For scalar variables, ensure shape is (batch, 1) for batched outputs
                if isinstance(var, TorchScalarVariable):
                    if output.dim() == 1:
                        # Shape is (batch,), reshape to (batch, 1)
                        parsed_outputs[output_name] = output.unsqueeze(-1)
                    elif output.dim() == 0:
                        # Scalar output, reshape to (1, 1)
                        parsed_outputs[output_name] = output.unsqueeze(0).unsqueeze(-1)
                    else:
                        # Already has proper dimensions or higher, ensure last dim is 1
                        parsed_outputs[output_name] = (
                            output.squeeze().unsqueeze(-1)
                            if output.squeeze().dim() > 0
                            else output.squeeze().unsqueeze(0).unsqueeze(-1)
                        )
                else:
                    # For non-scalar outputs (NDVariable), keep original behavior
                    parsed_outputs[output_name] = output.squeeze()

        return parsed_outputs

    def _prepare_outputs(
        self,
        parsed_outputs: dict[str, torch.Tensor],
    ) -> dict[str, Union[float, torch.Tensor]]:
        """Update and return outputs according to ``output_format``.

        Updates the output variables within the model to reflect the new
        values.

        Parameters
        ----------
        parsed_outputs : dict of str to torch.Tensor
            Dictionary of output variable names to transformed tensors.

        Returns
        -------
        dict of str to float or torch.Tensor
            Dictionary of output variable names to values depending on
            ``output_format``.

        """
        if self.output_format.lower() == "tensor":
            return parsed_outputs
        else:
            return {
                key: value.item() if value.squeeze().dim() == 0 else value
                for key, value in parsed_outputs.items()
            }

__init__(*args, **kwargs)

Initializes TorchModel.

Parameters

args : dict, str, or Path Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path. *kwargs See class attributes.

Source code in lume_torch/models/torch_model.py

def __init__(self, *args, **kwargs):
    """Initializes TorchModel.

    Parameters
    ----------
    *args : dict, str, or Path
        Accepts a single argument which is the model configuration as dictionary, YAML or JSON
        formatted string or file path.
    **kwargs
        See class attributes.

    """
    super().__init__(*args, **kwargs)

    # dtype property sets precision across model and transformers
    self.dtype

    # fixed model: set full model in eval mode and deactivate all gradients
    if self.fixed_model:
        is_scripted = isinstance(self.model, torch.jit.ScriptModule)
        self.model.eval().requires_grad_(False) if not is_scripted else None
        for t in self.input_transformers + self.output_transformers:
            if isinstance(t, torch.nn.Module):
                t.eval().requires_grad_(False)

    # ensure consistent device
    self.to(self.device)

_evaluate(input_dict)

Evaluate the model on the given input dictionary.

Parameters

input_dict : dict of str to float or torch.Tensor Input dictionary on which to evaluate the model.

Returns

dict of str to float or torch.Tensor Dictionary of output variable names to values.

Source code in lume_torch/models/torch_model.py

def _evaluate(
    self,
    input_dict: dict[str, Union[float, torch.Tensor]],
) -> dict[str, Union[float, torch.Tensor]]:
    """Evaluate the model on the given input dictionary.

    Parameters
    ----------
    input_dict : dict of str to float or torch.Tensor
        Input dictionary on which to evaluate the model.

    Returns
    -------
    dict of str to float or torch.Tensor
        Dictionary of output variable names to values.

    """
    formatted_inputs = format_inputs(input_dict)
    input_tensor = self._arrange_inputs(formatted_inputs)
    input_tensor = self._transform_inputs(input_tensor)
    output_tensor = self.model(input_tensor)
    output_tensor = self._transform_outputs(output_tensor)
    parsed_outputs = self._parse_outputs(output_tensor)
    output_dict = self._prepare_outputs(parsed_outputs)
    return output_dict

Usage Example

from lume_torch.models.torch_model import TorchModel
from lume_torch.variables import ScalarVariable
import torch.nn as nn

# Create a neural network
network = nn.Sequential(
    nn.Linear(2, 10),
    nn.ReLU(),
    nn.Linear(10, 1)
)

# Wrap in TorchModel
model = TorchModel(
    model=network,
    input_variables=[
        ScalarVariable(name="x1", value_range=[-5, 5]),
        ScalarVariable(name="x2", value_range=[-5, 5]),
    ],
    output_variables=[
        ScalarVariable(name="y"),
    ],
    device="cpu",
    precision="double"
)

# Evaluate
result = model({"x1": 1.0, "x2": 2.0})

TorchModule

PyTorch-compatible interface for LUME models.

lume_torch.models.torch_module.TorchModule

Bases: Module

Wrapper to allow a LUME TorchModel to be used like a torch.nn.Module.

As the base model within the TorchModel is assumed to be fixed during instantiation, so is the TorchModule.

Source code in lume_torch/models/torch_module.py

class TorchModule(torch.nn.Module):
    """Wrapper to allow a LUME TorchModel to be used like a torch.nn.Module.

    As the base model within the TorchModel is assumed to be fixed during instantiation,
    so is the TorchModule.

    """

    def __init__(
        self,
        *args,
        model: TorchModel = None,
        input_order: list[str] = None,
        output_order: list[str] = None,
    ):
        """Initializes TorchModule.

        Parameters
        ----------
        *args : dict, str, or Path
            Accepts a single argument which is the model configuration as dictionary, YAML or JSON
            formatted string or file path.
        model : TorchModel, optional
            The TorchModel instance to wrap around. If config is None, this has to be defined.
        input_order : list of str, optional
            Input names in the order they are passed to the model. If None, the input order of the
            TorchModel is used.
        output_order : list of str, optional
            Output names in the order they are returned by the model. If None, the output order of
            the TorchModel is used.

        """
        if all(arg is None for arg in [*args, model]):
            logger.error("TorchModule requires either a YAML config or model argument")
            raise ValueError(
                "Either a YAML string has to be given or model has to be defined."
            )
        super().__init__()
        if len(args) == 1:
            if not all(v is None for v in [model, input_order, output_order]):
                logger.error(
                    "Cannot specify both YAML config and keyword arguments for TorchModule"
                )
                raise ValueError(
                    "Cannot specify YAML string and keyword arguments for TorchModule init."
                )
            logger.debug("Initializing TorchModule from configuration file")
            model_fields = {f"model.{k}": v for k, v in TorchModel.model_fields.items()}
            kwargs = parse_config(args[0], model_fields)
            kwargs["model"] = TorchModel(kwargs["model"])
            self.__init__(**kwargs)
        elif len(args) > 1:
            logger.error(f"Too many positional arguments to TorchModule: {len(args)}")
            raise ValueError(
                "Arguments to TorchModule must be either a single YAML string or keyword arguments."
            )
        else:
            logger.debug(f"Initializing TorchModule with model: {type(model).__name__}")
            self._model = model
            self._input_order = input_order
            self._output_order = output_order
            self.register_module("base_model", self._model.model)
            logger.debug(
                f"Registered {len(self._model.input_transformers)} input transformers"
            )
            for i, input_transformer in enumerate(self._model.input_transformers):
                self.register_module(f"input_transformers_{i}", input_transformer)
            logger.debug(
                f"Registered {len(self._model.output_transformers)} output transformers"
            )
            for i, output_transformer in enumerate(self._model.output_transformers):
                self.register_module(f"output_transformers_{i}", output_transformer)
            if not model.model.training:  # TorchModel defines train/eval mode
                self.eval()
            logger.info(
                f"Initialized TorchModule with {len(self.input_order)} inputs and {len(self.output_order)} outputs"
            )

    @property
    def model(self):
        return self._model

    @property
    def input_order(self):
        if self._input_order is None:
            return self._model.input_names
        else:
            return self._input_order

    @property
    def output_order(self):
        if self._output_order is None:
            return self._model.output_names
        else:
            return self._output_order

    def forward(self, x: torch.Tensor):
        # input shape: [..., n_features] or [..., n_features, 1] for scalar variables
        x = self._validate_input(x)
        model_input = self._tensor_to_dictionary(x)
        y_model = self.evaluate_model(model_input)
        y_model = self.manipulate_output(y_model)
        # squeeze for use as prior mean in botorch GPs
        y = self._dictionary_to_tensor(y_model).squeeze()
        return y

    def yaml(
        self,
        base_key: str = "",
        file_prefix: str = "",
        save_models: bool = False,
        save_jit: bool = False,
    ) -> str:
        """Serializes the object and returns a YAML formatted string defining the TorchModule instance.

        Parameters
        ----------
        base_key : str, optional
            Base key for serialization.
        file_prefix : str, optional
            Prefix for generated filenames.
        save_models : bool, optional
            Determines whether models are saved to file.
        save_jit : bool, optional
            Determines whether the structure of the model is saved as TorchScript

        Returns
        -------
        str
            YAML formatted string defining the TorchModule instance.

        """
        d = {}
        for k, v in inspect.signature(TorchModule.__init__).parameters.items():
            if k not in ["self", "args", "model"]:
                d[k] = getattr(self, k)
        output = json.loads(
            json.dumps(
                recursive_serialize(d, base_key, file_prefix, save_models, save_jit)
            )
        )
        model_output = json.loads(
            self._model.to_json(
                base_key=base_key,
                file_prefix=file_prefix,
                save_models=save_models,
                save_jit=save_jit,
            )
        )
        output["model"] = model_output
        # create YAML formatted string
        s = yaml.dump(
            {"model_class": self.__class__.__name__} | output,
            default_flow_style=None,
            sort_keys=False,
        )
        return s

    def dump(
        self,
        file: Union[str, os.PathLike],
        save_models: bool = True,
        base_key: str = "",
        save_jit: bool = False,
    ):
        """Returns and optionally saves YAML formatted string defining the model.

        Parameters
        ----------
        file : str or Path
            File path to which the YAML formatted string and corresponding files are saved.
        save_models : bool, optional
            Determines whether models are saved to file.
        base_key : str, optional
            Base key for serialization.
        save_jit : bool, optional
            Whether the model is saved using just in time pytorch method

        """
        logger.info(f"Dumping TorchModule configuration to: {file}")
        if save_models:
            logger.debug("Saving model files alongside configuration")
        if save_jit:
            logger.debug("Saving TorchModule as TorchScript (JIT)")
        file_prefix = os.path.splitext(file)[0]
        with open(file, "w") as f:
            f.write(
                self.yaml(
                    save_models=save_models,
                    base_key=base_key,
                    file_prefix=file_prefix,
                    save_jit=save_jit,
                )
            )

    def evaluate_model(self, x: dict[str, torch.Tensor]):
        """Placeholder method to modify model calls.

        Parameters
        ----------
        x : dict of str to torch.Tensor
            Input dictionary to evaluate.

        Returns
        -------
        dict of str to torch.Tensor
            Model evaluation results.

        """
        return self._model.evaluate(x)

    def manipulate_output(self, y_model: dict[str, torch.Tensor]):
        """Placeholder method to modify the model output.

        Parameters
        ----------
        y_model : dict of str to torch.Tensor
            Model output dictionary.

        Returns
        -------
        dict of str to torch.Tensor
            Modified model output.

        """
        return y_model

    def _tensor_to_dictionary(self, x: torch.Tensor):
        input_dict = {}
        # Handle both old format (..., n_features) and new format (..., n_features, 1)
        # New format requires at least 3 dimensions with last dim == 1
        if x.ndim >= 3 and x.shape[-1] == 1:
            # New scalar format: (..., n_features, 1)
            # Index the second-to-last dimension and keep trailing 1
            for idx, input_name in enumerate(self.input_order):
                input_dict[input_name] = x[..., idx, :]
        else:
            # Old format: (..., n_features)
            # Index last dimension and add trailing 1
            for idx, input_name in enumerate(self.input_order):
                input_dict[input_name] = x[..., idx].unsqueeze(-1)
        return input_dict

    def _dictionary_to_tensor(self, y_model: dict[str, torch.Tensor]):
        # Model outputs have shape (batch, 1) for single-sample or (batch, samples) for multi-sample
        # We need to stack them into (..., n_outputs) format
        output_list = []
        for output_name in self.output_order:
            output = y_model[output_name]
            # If output has trailing 1 (single-sample), squeeze it for stacking
            # Multi-sample outputs like (batch, samples) are kept as-is
            if output.shape[-1] == 1:
                output = output.squeeze(-1)
            output_list.append(output)

        output_tensor = torch.stack(output_list, dim=-1)
        return output_tensor

    @staticmethod
    def _validate_input(x: torch.Tensor) -> torch.Tensor:
        if x.dim() <= 1:
            logger.error(
                f"Invalid input dimensions: expected at least 2D ([n_samples, n_features]), got {tuple(x.shape)}"
            )
            raise ValueError(
                f"Expected input dim to be at least 2 ([n_samples, n_features]), received: {tuple(x.shape)}"
            )
        else:
            return x

    def register_to_mlflow(
        self,
        artifact_path: str,
        registered_model_name: str | None = None,
        tags: dict[str, Any] | None = None,
        version_tags: dict[str, Any] | None = None,
        alias: str | None = None,
        run_name: str | None = None,
        log_model_dump: bool = True,
        save_jit: bool = False,
        **kwargs,
    ):
        """Registers the model to MLflow if mlflow is installed.

        Each time this function is called, a new version of the model is created. The model is saved to the
        tracking server or local directory, depending on the MLFLOW_TRACKING_URI.

        If no tracking server is set up, data and artifacts are saved directly under your current directory. To set up
        a tracking server, set the environment variable MLFLOW_TRACKING_URI, e.g. a local port/path. See
        https://mlflow.org/docs/latest/getting-started/intro-quickstart/ for more info.

        Parameters
        ----------
        artifact_path : str
            Path to store the model in MLflow.
        registered_model_name : str or None, optional
            Name of the registered model in MLflow.
        tags : dict of str to Any or None, optional
            Tags to add to the MLflow model.
        version_tags : dict of str to Any or None, optional
            Tags to add to this MLflow model version.
        alias : str or None, optional
            Alias to add to this MLflow model version.
        run_name : str or None, optional
            Name of the MLflow run.
        log_model_dump : bool, optional
            Whether to log the model dump files as artifacts.
        save_jit : bool, optional
            Whether to save the model as TorchScript when calling model.dump, if log_model_dump=True.
        **kwargs
            Additional arguments for mlflow.pyfunc.log_model.

        Returns
        -------
        mlflow.models.model.ModelInfo
            Model info metadata.

        """
        return register_model(
            self,
            artifact_path,
            registered_model_name,
            tags,
            version_tags,
            alias,
            run_name,
            log_model_dump,
            save_jit,
            **kwargs,
        )

__init__(*args, model=None, input_order=None, output_order=None)

Initializes TorchModule.

Parameters

*args : dict, str, or Path Accepts a single argument which is the model configuration as dictionary, YAML or JSON formatted string or file path. model : TorchModel, optional The TorchModel instance to wrap around. If config is None, this has to be defined. input_order : list of str, optional Input names in the order they are passed to the model. If None, the input order of the TorchModel is used. output_order : list of str, optional Output names in the order they are returned by the model. If None, the output order of the TorchModel is used.

Source code in lume_torch/models/torch_module.py

def __init__(
    self,
    *args,
    model: TorchModel = None,
    input_order: list[str] = None,
    output_order: list[str] = None,
):
    """Initializes TorchModule.

    Parameters
    ----------
    *args : dict, str, or Path
        Accepts a single argument which is the model configuration as dictionary, YAML or JSON
        formatted string or file path.
    model : TorchModel, optional
        The TorchModel instance to wrap around. If config is None, this has to be defined.
    input_order : list of str, optional
        Input names in the order they are passed to the model. If None, the input order of the
        TorchModel is used.
    output_order : list of str, optional
        Output names in the order they are returned by the model. If None, the output order of
        the TorchModel is used.

    """
    if all(arg is None for arg in [*args, model]):
        logger.error("TorchModule requires either a YAML config or model argument")
        raise ValueError(
            "Either a YAML string has to be given or model has to be defined."
        )
    super().__init__()
    if len(args) == 1:
        if not all(v is None for v in [model, input_order, output_order]):
            logger.error(
                "Cannot specify both YAML config and keyword arguments for TorchModule"
            )
            raise ValueError(
                "Cannot specify YAML string and keyword arguments for TorchModule init."
            )
        logger.debug("Initializing TorchModule from configuration file")
        model_fields = {f"model.{k}": v for k, v in TorchModel.model_fields.items()}
        kwargs = parse_config(args[0], model_fields)
        kwargs["model"] = TorchModel(kwargs["model"])
        self.__init__(**kwargs)
    elif len(args) > 1:
        logger.error(f"Too many positional arguments to TorchModule: {len(args)}")
        raise ValueError(
            "Arguments to TorchModule must be either a single YAML string or keyword arguments."
        )
    else:
        logger.debug(f"Initializing TorchModule with model: {type(model).__name__}")
        self._model = model
        self._input_order = input_order
        self._output_order = output_order
        self.register_module("base_model", self._model.model)
        logger.debug(
            f"Registered {len(self._model.input_transformers)} input transformers"
        )
        for i, input_transformer in enumerate(self._model.input_transformers):
            self.register_module(f"input_transformers_{i}", input_transformer)
        logger.debug(
            f"Registered {len(self._model.output_transformers)} output transformers"
        )
        for i, output_transformer in enumerate(self._model.output_transformers):
            self.register_module(f"output_transformers_{i}", output_transformer)
        if not model.model.training:  # TorchModel defines train/eval mode
            self.eval()
        logger.info(
            f"Initialized TorchModule with {len(self.input_order)} inputs and {len(self.output_order)} outputs"
        )

forward(x)

Source code in lume_torch/models/torch_module.py

def forward(self, x: torch.Tensor):
    # input shape: [..., n_features] or [..., n_features, 1] for scalar variables
    x = self._validate_input(x)
    model_input = self._tensor_to_dictionary(x)
    y_model = self.evaluate_model(model_input)
    y_model = self.manipulate_output(y_model)
    # squeeze for use as prior mean in botorch GPs
    y = self._dictionary_to_tensor(y_model).squeeze()
    return y

Usage Example

from lume_torch.models.torch_module import TorchModule
import torch

# Wrap TorchModel in TorchModule
torch_module = TorchModule(model=model)

# Use like a PyTorch module
input_tensor = torch.tensor([[1.0, 2.0]])
output_tensor = torch_module(input_tensor)

# Integrate with PyTorch pipelines
optimizer = torch.optim.Adam(torch_module.parameters())

Probabilistic Models

For models that output distributions, see Probabilistic Models.

Serialization Functions

lume_torch.base.process_torch_module(module, base_key='', key='', file_prefix='', save_modules=True, save_jit=False)

Optionally saves the given torch module to file and returns the filename.

Parameters

module : torch.nn.Module The torch module to process. base_key : str, optional Base key at this stage of serialization. key : str, optional Key corresponding to the torch module. file_prefix : str or os.PathLike, optional Prefix for generated filenames. save_modules : bool, optional Determines whether torch modules are saved to file. save_jit : bool, optional Determines whether the model gets saved as TorchScript.

Returns

str Filename under which the torch module is (or would be) saved.

Source code in lume_torch/base.py

def process_torch_module(
    module,
    base_key: str = "",
    key: str = "",
    file_prefix: Union[str, os.PathLike] = "",
    save_modules: bool = True,
    save_jit: bool = False,
):
    """Optionally saves the given torch module to file and returns the filename.

    Parameters
    ----------
    module : torch.nn.Module
        The torch module to process.
    base_key : str, optional
        Base key at this stage of serialization.
    key : str, optional
        Key corresponding to the torch module.
    file_prefix : str or os.PathLike, optional
        Prefix for generated filenames.
    save_modules : bool, optional
        Determines whether torch modules are saved to file.
    save_jit : bool, optional
        Determines whether the model gets saved as TorchScript.

    Returns
    -------
    str
        Filename under which the torch module is (or would be) saved.

    """
    torch = try_import_module("torch")
    filepath_prefix, filename_prefix = os.path.split(file_prefix)
    prefixes = [ele for ele in [filename_prefix, base_key] if not ele == ""]
    filename = "{}.pt".format(key)
    jit_filename = "{}.jit".format(key)
    if prefixes:
        filename = "_".join((*prefixes, filename))
        jit_filename = "_".join((*prefixes, jit_filename))
    if save_modules:
        filepath = os.path.join(filepath_prefix, filename)
        torch.save(module, filepath)
        logger.debug(f"Saved torch module to: {filepath}")
    if save_jit:
        filepath = os.path.join(filepath_prefix, jit_filename)
        try:
            scripted_model = torch.jit.script(module)
            torch.jit.save(scripted_model, filepath)
            logger.debug(f"Saved JIT model to: {filepath}")
        except Exception as e:
            logger.warning(
                "Saving as JIT through scripting has only been evaluated "
                "for NN models that don't depend on BoTorch modules."
            )
            logger.error(f"Failed to script the model: {e}")
            raise e
    return jit_filename if save_jit else filename

lume_torch.base.model_kwargs_from_dict(config)

Processes model configuration and returns the corresponding keyword arguments for model constructor.

Parameters

config : dict Model configuration.

Returns

dict Configuration as keyword arguments for model constructor.

Source code in lume_torch/base.py

def model_kwargs_from_dict(config: dict) -> dict:
    """Processes model configuration and returns the corresponding keyword arguments for model constructor.

    Parameters
    ----------
    config : dict
        Model configuration.

    Returns
    -------
    dict
        Configuration as keyword arguments for model constructor.

    """
    config = deserialize_variables(config)
    if all(key in config.keys() for key in ["input_variables", "output_variables"]):
        config["input_variables"], config["output_variables"] = variables_from_dict(
            config
        )
    config.pop("model_class", None)
    return config

lume_torch.base.parse_config(config, model_fields=None)

Parses model configuration and returns keyword arguments for model constructor.

Parameters

config : dict, str, TextIOWrapper, or os.PathLike Model configuration as dictionary, YAML or JSON formatted string, file or file path. model_fields : dict, optional Fields expected by the model (required for replacing relative paths).

Returns

dict Configuration as keyword arguments for model constructor.

Source code in lume_torch/base.py

def parse_config(
    config: Union[dict, str, TextIOWrapper, os.PathLike],
    model_fields: dict = None,
) -> dict:
    """Parses model configuration and returns keyword arguments for model constructor.

    Parameters
    ----------
    config : dict, str, TextIOWrapper, or os.PathLike
        Model configuration as dictionary, YAML or JSON formatted string, file or file path.
    model_fields : dict, optional
        Fields expected by the model (required for replacing relative paths).

    Returns
    -------
    dict
        Configuration as keyword arguments for model constructor.

    """
    config_file = None
    if isinstance(config, dict):
        logger.debug("Parsing configuration from dictionary")
        d = config
    else:
        if isinstance(config, TextIOWrapper):
            logger.debug(f"Reading configuration from file wrapper: {config.name}")
            yaml_str = config.read()
            config_file = os.path.abspath(config.name)
        elif isinstance(config, (str, os.PathLike)) and os.path.exists(config):
            logger.debug(f"Loading configuration from file: {config}")
            with open(config) as f:
                yaml_str = f.read()
            config_file = os.path.abspath(config)
        else:
            logger.debug("Parsing configuration from YAML string")
            yaml_str = config
        d = recursive_deserialize(yaml.safe_load(yaml_str))
    if config_file is not None:
        config_dir = os.path.dirname(os.path.realpath(config_file))
        logger.debug(f"Replacing relative paths using config directory: {config_dir}")
        d = replace_relative_paths(d, model_fields, config_dir)
    return model_kwargs_from_dict(d)

lume_torch.base.recursive_serialize(v, base_key='', file_prefix='', save_models=True, save_jit=False)

Recursively performs custom serialization for the given object.

Parameters

v : dict of str to Any Object to serialize. base_key : str, optional Base key at this stage of serialization. file_prefix : str or os.PathLike, optional Prefix for generated filenames. save_models : bool, optional Determines whether models are saved to file. save_jit : bool, optional Determines whether the model is saved as TorchScript.

Returns

dict Serialized object.

Source code in lume_torch/base.py

def recursive_serialize(
    v: dict[str, Any],
    base_key: str = "",
    file_prefix: Union[str, os.PathLike] = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Recursively performs custom serialization for the given object.

    Parameters
    ----------
    v : dict of str to Any
        Object to serialize.
    base_key : str, optional
        Base key at this stage of serialization.
    file_prefix : str or os.PathLike, optional
        Prefix for generated filenames.
    save_models : bool, optional
        Determines whether models are saved to file.
    save_jit : bool, optional
        Determines whether the model is saved as TorchScript.

    Returns
    -------
    dict
        Serialized object.

    """
    logger.debug(
        f"Serializing object with base_key: '{base_key}', {len(v)} top-level keys"
    )
    # try to import modules for LUMETorch child classes
    torch = try_import_module("torch")
    # serialize
    v = serialize_variables(v)
    for key, value in v.items():
        if isinstance(value, dict):
            logger.debug(f"Recursively serializing nested dict for key: '{key}'")
            v[key] = recursive_serialize(value, key)
        elif isinstance(value, list) and all(isinstance(ele, dict) for ele in value):
            # e.g. NN ensemble
            logger.debug(
                f"Serializing NN ensemble with {len(value)} models for key: '{key}'"
            )
            v[key] = [
                recursive_serialize(
                    value[i], f"{base_key}{i}", file_prefix, save_models, save_jit
                )
                for i in range(len(value))
            ]
            # For NN ensembles, we want v[key] to be a list of the filenames corresponding to each
            # model in the ensemble and not the serialized dict of each
            # NOTE: If this clause is reached for other models, we may need to do this differently
            v[key] = [v[key][i]["model"] for i in range(len(value))]
        elif torch is not None and isinstance(value, torch.nn.Module):
            logger.debug(f"Serializing torch.nn.Module for key: '{key}'")
            v[key] = process_torch_module(
                value, base_key, key, file_prefix, save_models, save_jit
            )
        elif (
            isinstance(value, list)
            and torch is not None
            and any(isinstance(ele, torch.nn.Module) for ele in value)
        ):
            # List of transformers
            logger.debug(
                f"Serializing {len(value)} torch.nn.Module transformers for key: '{key}'"
            )
            v[key] = [
                process_torch_module(
                    value[i], base_key, f"{key}_{i}", file_prefix, save_models, False
                )
                for i in range(len(value))
            ]
        else:
            for _type, func in JSON_ENCODERS.items():
                if isinstance(value, _type):
                    logger.debug(
                        f"Applying JSON encoder for type {_type.__name__} to key: '{key}'"
                    )
                    v[key] = func(value)
        # check to make sure object has been serialized, if not use a generic serializer
        try:
            json.dumps(v[key])
        except (TypeError, OverflowError):
            logger.debug(
                f"Using generic serializer for unserializable object at key: '{key}'"
            )
            v[key] = f"{v[key].__module__}.{v[key].__class__.__qualname__}"

    return v

lume_torch.base.recursive_deserialize(v)

Recursively performs custom deserialization for the given object.

Parameters

v : dict Object to deserialize.

Returns

dict Deserialized object.

Source code in lume_torch/base.py

def recursive_deserialize(v):
    """Recursively performs custom deserialization for the given object.

    Parameters
    ----------
    v : dict
        Object to deserialize.

    Returns
    -------
    dict
        Deserialized object.

    """
    logger.debug(f"Deserializing object with {len(v)} top-level keys")
    # deserialize
    v = deserialize_variables(v)
    for key, value in v.items():
        if isinstance(value, dict):
            logger.debug(f"Recursively deserializing nested dict for key: '{key}'")
            v[key] = recursive_deserialize(value)
    return v

lume_torch.base.json_dumps(v, *, base_key='', file_prefix='', save_models=True, save_jit=False)

Serializes variables before dumping with json.

Parameters

v : object Object to dump. base_key : str, optional Base key for serialization. file_prefix : str or os.PathLike, optional Prefix for generated filenames. save_models : bool, optional Determines whether models are saved to file. save_jit : bool, optional Determines whether the model is saved as TorchScript.

Returns

str JSON formatted string.

Source code in lume_torch/base.py

def json_dumps(
    v,
    *,
    base_key="",
    file_prefix: Union[str, os.PathLike] = "",
    save_models: bool = True,
    save_jit: bool = False,
):
    """Serializes variables before dumping with json.

    Parameters
    ----------
    v : object
        Object to dump.
    base_key : str, optional
        Base key for serialization.
    file_prefix : str or os.PathLike, optional
        Prefix for generated filenames.
    save_models : bool, optional
        Determines whether models are saved to file.
    save_jit : bool, optional
        Determines whether the model is saved as TorchScript.

    Returns
    -------
    str
        JSON formatted string.

    """
    v = recursive_serialize(
        v.model_dump(), base_key, file_prefix, save_models, save_jit
    )
    v = json.dumps(v)
    return v

lume_torch.base.json_loads(v)

Loads JSON formatted string and recursively deserializes the result.

Parameters

v : str JSON formatted string to load.

Returns

dict Deserialized object.

Source code in lume_torch/base.py

def json_loads(v):
    """Loads JSON formatted string and recursively deserializes the result.

    Parameters
    ----------
    v : str
        JSON formatted string to load.

    Returns
    -------
    dict
        Deserialized object.

    """
    v = json.loads(v)
    v = recursive_deserialize(v)
    return v

Configuration Files

Models can be saved and loaded using YAML configuration files:

# Save model
model.dump("my_model.yml")

# Load model
loaded_model = MyModel("my_model.yml")

Example configuration file:

model_class: TorchModel
input_variables:
  x1:
    variable_class: ScalarVariable
    default_value: 0.0
    value_range: [-5.0, 5.0]
  x2:
    variable_class: ScalarVariable
    default_value: 0.0
    value_range: [-5.0, 5.0]
output_variables:
  y:
    variable_class: ScalarVariable
model: model.pt
device: cpu
precision: double

See Also

• Variables - Defining inputs and outputs
• Probabilistic Models - GP and ensemble models
• Examples - Example notebooks