Utilities Reference

This page contains the API reference for all utility functions and classes throughout the library.

---

PyTorch Weight Conversion

These utilities help import pre-trained weights from PyTorch models (specifically the neuraloperator library) into their native JAX neojax equivalents.

neojax.utils.torch_converters.load_torch_weights_into_fno(model: FNO, state_dict: dict[str, typing.Any]) -> FNO

Loads weights from a PyTorch FNO model state dict into a Neojax FNO instance.

Handles necessary shape transpositions (e.g. PyTorch [in, out, ...] vs JAX [out, in, ...] for spectral weights).

Parameters:

Name	Type	Description	Default
model	FNO	The target Neojax FNO model.	required
state_dict	dict[str, Any]	PyTorch FNO model state dict.	required

Returns:

Type	Description
FNO	A new FNO model instance with loaded weights.

neojax.utils.torch_converters.load_torch_weights_into_geo_fno(model: GeoFNO, state_dict: dict[str, typing.Any]) -> GeoFNO

Loads weights from a PyTorch Geo-FNO model state dict into a Neojax GeoFNO instance.

Handles coordinate diffeomorphism (geomap) parameters and intermediate coordinate biases.

Parameters:

Name	Type	Description	Default
model	GeoFNO	The target Neojax GeoFNO model.	required
state_dict	dict[str, Any]	PyTorch Geo-FNO model state dict.	required

Returns:

Type	Description
GeoFNO	A new GeoFNO model instance with loaded weights.

---

Other Utilites

neojax.utils.tree_utils.tree_add_broadcast_to(tree: PyTree, *, axis: int, size: int, is_leaf: collections.abc.Callable[[typing.Any], bool] | None = None) -> PyTree

Broadcasts the given axis of each leaf of the given tree to given size.

How this function works depends on if the given axis has the given size or not:

• If the given axis already has the given size, the leaf is returned unchanged.
• If the given axis has size 1, this axis is broadcasted to size.
• If the given axis is not broadcastable to size, a new axis is inserted and broadcasted.

Parameters:

Name	Type	Description	Default
tree	PyTree	Pytree of arrays to broadcast.	required
axis	int	Array axis to add.	required
size	int	Size of given axis.	required
is_leaf	Callable[[Any], bool] | None	Optional callable specifying whether to traverse the current object at each flattening step. Default is None, that is only arrays are treated as leaves.	None

Returns:

Type	Description
PyTree	Pytree of arrays with additional axis of size length.

Examples:

>>> import jax.numpy as jnp
>>> pytree = {"one": jnp.ones((4, 2, 1, 3)), "two": jnp.ones((4, 2, 3)) * 2.0}
>>> axis = 2
>>> size = 3
>>> bc_tree = tree_add_broadcast_to(pytree, axis=axis, size=size)
>>> bc_tree["one"].shape
(4, 2, 3, 3)
>>> bc_tree["two"].shape
(4, 2, 3)

Out-of-bounds axis

This function silently wraps an out-of-bounds axis via modulo.