nifty.re.model module

class ClipModel(model: Model, threshold: float = 10.0, warn: bool = False, custom_clip_func=None)

Bases: Model

A wrapper around a NIFTy model that clips all input values to a specified threshold before passing them to the underlying model.

This is useful for preventing numerical instabilities caused by extreme values in latent variables. However, this model is mostly intended for testing and debugging rather than production codes.

Parameters:

• model (Model) – The NIFTy model to be wrapped. This model is called on the clipped version of the input.

• threshold (float, default=10.0) – The absolute value used for default clipping. When custom_clip_func is not provided, every leaf array in the input pytree is clipped elementwise to the interval [-threshold, threshold].

• warn (bool, default=False) – If True, a warning is emitted whenever any element in the input pytree exceeds threshold in absolute value prior to clipping.

• custom_clip_func (callable, optional) – A custom function applied to each leaf of the input pytree instead of jnp.clip. It should take a single JAX array and return a transformed array. If provided, threshold is not used for clipping, but is still used for the warning check.

__init__(model: Model, threshold: float = 10.0, warn: bool = False, custom_clip_func=None)

clip: Callable

model: Model

threshold: float

warn: bool

class Initializer(call_or_struct)

Bases: object

domain = ShapeWithDtype(shape=(2,), dtype=<class 'jax.numpy.uint32'>)

property stupid

property target

class LazyModel(domain=<class 'nifty.re.model.NoValue'>, target=<class 'nifty.re.model.NoValue'>, init=<class 'nifty.re.model.NoValue'>)

Bases: object

Base model with lazy evaluation of domain, target, and initializer.

Properties automatically derive:

• domain from init via eval_shape (if not provided)

• target from __call__ and domain via eval_shape (if not provided)

• A default white-noise initializer (if not provided)

See ModelMeta for details on JAX PyTree registration and static vs. dynamic fields.

__init__(domain=<class 'nifty.re.model.NoValue'>, target=<class 'nifty.re.model.NoValue'>, init=<class 'nifty.re.model.NoValue'>)

property domain

property init: Initializer

property target

class Model(call: ~typing.Callable | None = None, *, domain=<class 'nifty.re.model.NoValue'>, target=<class 'nifty.re.model.NoValue'>, init=<class 'nifty.re.model.NoValue'>, white_init=False)

Bases: LazyModel

Main building block for Nifty.re models.

Joins a callable with domain, target, and initializer method. From a domain and callable, automatically derives the target and instantiates a default initializer if not set explicitly.

This class is wrapped in a dataclass and registered as a JAX PyTree, allowing it to be used with transformations such as jit, vmap, or grad. By default, all fields are marked as static (compile-time constants). To prevent large arrays from being inlined into compiled code, mark them as dynamic in the class definition:

my_array: Any = dataclasses.field(metadata=dict(static=False))

Parameters:

• call (callable) – Method acting on objects of type domain.

• domain (tree-like structure of ShapeWithDtype, optional) – PyTree of objects with a shape and dtype attribute. Inferred from init if not specified.

• target (tree-like structure of ShapeWithDtype, optional) – PyTree of objects with a shape and dtype attribute akin to the output of call. Inferred from call and domain if not set.

• init (callable, optional) – Initialization method taking a PRNG key as first argument and creating an object of type domain. Inferred from domain assuming a white normal prior if not set.

• white_init (bool, optional) – If True, the domain is set to a white normal prior. Defaults to False.

Notes

When composing models hierarchically, sub-models should be marked as dynamic:

class SubModel(jft.Model):
    ...

class ParentModel(jft.Model):
    sub: SubModel = dataclasses.field(metadata=dict(static=False))

Note that the static/dynamic classification within the sub-model is preserved: fields marked as static in the sub-model remain static, even though the sub-model itself is marked as dynamic in the parent model. The dynamic marking only ensures that JAX recursively flattens the sub-model’s PyTree structure instead of treating it as a single static leaf.

__init__(call: ~typing.Callable | None = None, *, domain=<class 'nifty.re.model.NoValue'>, target=<class 'nifty.re.model.NoValue'>, init=<class 'nifty.re.model.NoValue'>, white_init=False)

class ModelMeta(name, bases, dict_, /, **kwargs)

Bases: ABCMeta

Metaclass that registers derived classes as JAX PyTrees and wraps them in dataclasses.

Fields are classified as either static or dynamic:

• Static (default): Treated as compile-time constants. Suitable for e.g. configuration parameters or hyperparameters.

• Dynamic: Treated as runtime values. Required to prevent large arrays from being inlined into compiled code.

To mark a field as dynamic, use:

my_array : Any = dataclasses.field(metadata=dict(static=False))

class NoValue

Bases: object

class VModel(model, axis_size, in_axes=0, out_axes=0)

Bases: LazyModel

__init__(model, axis_size, in_axes=0, out_axes=0)

axis_size: int

in_axes: Any

model: LazyModel

out_axes: Any

class WrappedCall(call: ~typing.Callable, *, name=None, shape=(), dtype=None, white_init=False, target=<class 'nifty.re.model.NoValue'>)

Bases: Model

Model wrapper that selects a named subset from the input.

Transforms call so that instead of acting on input directly, it selects input[name] from a dict-like input structure.

See Model for details on the remaining arguments.

Parameters:

• call (callable) – Callable to wrap.

• name (hashable, optional) – Key used to select from input before passing to call.

• shape (tuple or tree-like structure of ShapeWithDtype) – Shape of old input on which call acts. This can also be an arbitrary shape-dtype structure in which case dtype is ignored. Defaults to a scalar.

• dtype (dtype or tree-like structure of ShapeWithDtype) – Dtype of the old input on which call acts. Redundant if shape already encodes the dtype.

__init__(call: ~typing.Callable, *, name=None, shape=(), dtype=None, white_init=False, target=<class 'nifty.re.model.NoValue'>)