Skip to content

BaseNO

Base Neural Operator¤

neojax.models.baseno.BaseNO ¤

Abstract base Neural Operator.

Defines shared functionalities for serialization, dtype conversion, inspection and profiling. All neojax models subclass BaseNO.

Custom Neural Operator classes should inherit from BaseNO.

save_weights(file: str | os.PathLike | typing.BinaryIO, **kwargs: typing.Any) -> None ¤

Saves only the model weights to the given file.

Parameters:

Name Type Description Default
file str | PathLike | BinaryIO

File descriptor or binary IO stream.

 required
kwargs Any

Additional arguments to pass to tree_serialise_leaves.

 required
load_weights(file: str | os.PathLike | typing.BinaryIO, **kwargs: typing.Any) -> BaseNO ¤

Loads only the model weights from the given file.

Parameters:

Name Type Description Default
file str | PathLike | BinaryIO

File descriptor or binary IO stream.

 required
kwargs Any

Additional arguments to pass to tree_deserialise_leaves.

 required
save(*, file: str | os.PathLike | typing.BinaryIO, hyperparams: dict[str, typing.Any], **kwargs: typing.Any) -> None ¤

Saves the model to the given file.

Parameters:

Name Type Description Default
file str | PathLike | BinaryIO

File descriptor or binary IO stream.

 required
hyperparams dict[str, Any]

The hyperparameters needed to replicate the model instance.

 required
kwargs Any

Additional arguments to pass to tree_serialise_leaves.

 required
load(*, file: str | os.PathLike | typing.BinaryIO, make_fn: collections.abc.Callable[[typing.Any], BaseNO], key: PRNGKeyArray | None = None, **kwargs: typing.Any) -> BaseNO classmethod ¤

Loads the model from the given file.

Parameters:

Name Type Description Default
file str | PathLike | BinaryIO

File descriptor or binary IO stream.

 required
make_fn Callable[[Any], BaseNO]

A function that takes all hyperparams saved to the model file and uses the necessary ones to construct the model and returns a model instance. See Equinox docs for details.

 required
key PRNGKeyArray | None

Optional random key passed to make_fn. If None and the model creation requires a random key, make sure to pass via functools.partial(make_fn, key=key). Default is None.

 None
kwargs Any

Additional arguments to pass to tree_deserialise_leaves.

 required

Returns:

Type Description
BaseNO

The loaded model instance.
astype(dtype: typing.Union[str, type[typing.Any], numpy.dtype, jax._src.typing.SupportsDType]) -> BaseNO ¤

Casts all inexact model parameters to the given dtype.

When casting to float, this method preserves the real/complex separation of complex parameters (by default, casting a complex to float in jax silently discards the imaginary part).

E.g. astype(jnp.float32) casts complex parameters to jnp.complex64 and real parameters to jnp.float32.

All real parameters are cast to the desired type.

Parameters:

Name Type Description Default
dtype Union[str, type[Any], dtype, SupportsDType]

Valid jaxtyping.DTypeLike to cast parameters to.

 required

Returns:

Type Description
BaseNO

A model instance with parameters cast to dtype.

Raises:

Type Description
ValueError

If an inadequate dtype is passed (e.g. any int).

Half float casting for complex arrays

Some Neural Operators (FNO, ...) have complex spectral weights. JAX doesn't support complex32, meaning casting a model with complex weights to half precision floats (jnp.float16, jnp.bfloat16) leaves the complex weights as jnp.complex64.

profile_compile(dummy_inputs: typing.Any, return_lowering: bool = False, filter_jax_frames: bool = True) -> tuple[bool, dict[str, typing.Any]] ¤

Profiles the XLA-compilation of the model.

This method relies on jax AOT lowering and compilation. See JAX docs for details.

Parameters:

Name Type Description Default
dummy_inputs Any

Dummy inputs representative of an input to model.call() Typically, a dummy array representative of a training sample to pass to the jitted model.

 required
return_lowering bool

Whether to return a human-readable representation of the model lowering in the summary dict.

 False
filter_jax_frames bool

In case of exceptions, whether to filter the jax and equinox frames to only focus on model code (jax internal traces are long).

 True

Returns:

Type Description
tuple[bool, dict[str, Any]]

True if the compilation succeeded else False and a summary dict of the compilation:

{

"lowering_cost_analysis": variable dict returned by `jax.stages.Lowered.cost_analysis()` or `None`,

"compiled_cost_analysis": variable dict returned by `jax.stages.Compiled.cost_analysis()` or `None`,

"exception": Possible exception string or `None`,

"lowering": Text representation of model lowering if `return_lowering` is `True` else empty string

}

Analysis Outputs

The outputs of the cost analysis vary with the underlying compiler. They may be None if the compiler doesn't provide such information.

size(return_n_params: bool = False) -> float | tuple[float, int] ¤

Computes model size in MB.

Parameters:

Name Type Description Default
return_n_params bool

Whether to return the number of parameters. Default is False.

 False

Returns:

Type Description
float | tuple[float, int]

The model size in MB and number of parameters if return_n_params is True.
__call__(x: Inexact[Array, 'in_c ...']) -> Inexact[Array, 'out_c ...'] ¤

Forward pass of the model.

Parameters:

Name Type Description Default
x Inexact[Array, 'in_c ...']

Input array shaped (in_c, d1, ..., dN).

 required

Returns:

Type Description
Inexact[Array, 'out_c ...']

Output array shaped (out_c, d1, ..., dN).

Custom Models¤

When implementing custom models, inherit from BaseNO to inherit useful functionality such as convenient saving and loading of models, etc.