global parameters vs. pure functions

[mkolodziejczyk-piap]

Hi, I'd like to ask what is the recommended way to use global parameters (constants) in functions. As stated in the docs:

JAX transformation and compilation are designed to work only on Python functions that are functionally pure: all the input data is passed through the function parameters, all the results are output through the function results. A pure function will always return the same result if invoked with the same inputs.

But usually we have a bunch of parameters, sometimes gathered in some kind of config object (e.g. dict). Do we have to pass it in each function to remain "pure"?

Here is an example from google mppi controller that use class constrants directly:

• class constants: https://github.com/google-research/google-research/blob/6b27191c314c6ed57437958bfabf5682c89b8bb5/jax_mpc/jax_mpc/mppi.py#L31-L32
• usage: https://github.com/google-research/google-research/blob/6b27191c314c6ed57437958bfabf5682c89b8bb5/jax_mpc/jax_mpc/mppi.py#L126-L127

BTW. Does "self" reference in class functions cause any problems?

Regards,

[jakevdp]

If you have global parameters that you can guarantee will stay constant, it's fine to reference them in compiled functions. For example:

from jax import jit

y = 1

@jit
def f(x):
  return x + y

print(f(1))
# 2

Such parameters are effectively treated as compile-time constants. Note however that if you mutate the parameters, it does not trigger a re-compile:

y = 10
print(f(1))
# 2

If you want mutations of your global state to be reflected in function calls, it's best to pass those parameters directly; e.g.

@jit
def f(x, y):
  return x + y

print(f(1, y))
# 11