Modular functions in JAX

[Mo-rpheus]

Hi JAX community!

I'm fairly new to JAX and really enjoying it so far (it's very accessible to someone from a non-maths/CS background). I have a software design related question to JAX.

To give a bit of context, I am trying to write a physical simulation package, with modularity as one of the features. For example, I want to have solvers that can be easily swapped out, and I am not sure what is the best way to go about this.

I know of one approach to do this:

from jax.tree_util import Partial

@Partial
def solver_a(...):
  pass

@Partial
def solver_b(...):
  pass

This works, but if the solver function needs any inputs unique to it compared to other solvers, then it quickly becomes unwieldy (closures, using Partial to freeze inputs, *args, etc. all can fix this, but it's a bit hacky imo).

I've been considering this approach recently:

from typing import NamedTuple

def solver_a(NamedTuple):
  unique_params: ArrayLike
  # any others unique params can go here

  # one implementation 
  def __call__(self, ...)->float:
    pass

def solver_b(NamedTuple):
  unique_params: float | int

  # another implementation
  def __call__(self, ...)->float:
    pass

The benefit of this approach is that named tuples are inherently immutable and transparent to JAX by default. Are there any downsides that I am not seeing?

Looking at similar packages (e.g. fluid simulations) written in JAX, I see a variety of different approaches. Some use dict, others register their classes as pytrees, and I've even seen some uses of equinox.Module. It's nice that JAX supports this many options, but what is best? What are the advantages/disadvantages of each approach? I would appreciate any thoughts on this :)

[jakevdp]

If you want to pass a function as a parameter, you can mark it static and that should be enough.

If you want to pass a more general structure, then you could use NamedTuple (if all elements in the structure are traceable objects like arrays) or more generally define a class with a custom pytree registration (equinox.Module allows for this, but is a pretty heavy dependency). For custom pytrees, the lightest-weight approach is to use register_dataclass; see https://jax.readthedocs.io/en/latest/_autosummary/jax.tree_util.register_dataclass.html for some examples (the easiest approach requires JAX v0.4.36, which should be released soon)

[patrick-kidger]

Unsurprisingly I'll advocate for equinox.Module as the best way to do this. :)

Of the options you list:

• dict: unstructured, probably not recommended in the same way that we prefer classes over dicts in regular Python.
• namedtuples: generally also recommended against in the same way as in regular Python; largely superseded by dataclasses.
• equinox.Module: this is just a dataclass-as-pytree, which handles all the edge cases for you (bound methods, inheritance, etc.) If you have multiple solvers that you want to treat interchangeably then I think one of the big advantages here is its strong support for abstract classes. (AbstractVar, __check_init__, see also this advanced guide.)
• custom classes-as-pytrees: similar to previous, doesn't have those features but drops the Equinox dependency.

Regardless of what you do, if you want interchangeable solvers then you might like to take a look at the solvers of libraries like Lineax, Optimistix, or Diffrax for inspiration, e.g. how lineax.linear_solve accepts a lineax.AbstractLinearSolver.