Best Practices for Composition of Configuration

[ntenenz]

Given the following pseudocode:

def g(x, a, b, c, d):
    # assume b, c, d must be static for jit but not a
    pass

def f(x1, a1, b1, c1, d1, x2, a2, b2, c2, d2, x3, a3, b3, c3, d3):
    out1 = g(x1, a1, b1, c1, d1)
    out2 = g(x2, a2, b2, c2, d2)
    out3 = g(x3, a3, b3, c3, d3)
    return out1 + out2 + out3

fjit = jax.jit(f, static_argnames=["b1", ...])
fjit(...)

The function signature of f there is already a bit extreme and can become even more so as g becomes more parameterized. One solution would be to construct a Config class:

def f(x1, x2, x3, config1, config2, config3):
    ...

fjit = jax.jit(f, static_argnames=["config1", "config2", "config3"]

In this case however, we still have these x1-x3 floating around and also cannot change a without triggering a recompile. Ideally, we could define a ProblemDef class containing x, a, b, c, and d and selectively label static only the elements which are required to be so. I don't believe this is currently possible with jit though.

Is there an alternative, established best practice for handling scenarios where some level of "static-ness" is required without resorting to flattened function signatures?

[YouJiacheng]

IIUC, custom Pytree exactly meets your requirements.

from typing import Hashable
from dataclasses import dataclass
import jax.numpy as jnp
from jax import tree_util

@dataclass
class ProblemDef:
    x: jnp.ndarray
    a: jnp.ndarray
    b: Hashable # Note: ndarray is not hashable, shim is needed
    c: Hashable
    d: Hashable
    
    def tree_flatten(self):
        children = (self.x, self.a) # arrays / dynamic values
        aux_data = (self.b, self.c, self.d) # hashable static values
        return children, aux_data
    
    @classmethod
    def tree_unflatten(cls, aux_data, children):
        x, a = children
        b, c, d = aux_data
        return cls(x, a, b, c, d)

tree_util.register_pytree_node_class(ProblemDef)
# Note: You should not mark ProblemDef as static argument

See also:
https://jax.readthedocs.io/en/latest/faq.html#strategy-3-making-customclass-a-pytree

[ntenenz]

I suspect I'm missing an obvious connection between jit/other function transformations and pytrees. The first two bullets represent my understanding of the framework, and the third is what I'm inferring from your comment. Could you please confirm that what I'm drawing from your statement is correct?

• Registering pytrees (specifying flattening/unflattening operations) allows one to extend the containers which jax is able to map, etc. as it enables jax to operate on flattened forms while providing the unflattened forms back to user code
• a jit trace can only follow a single code path and thus must be informed, explicitly (e.g., static_argnames) or implicitly (e.g., closure), of the static-ness of non-traceable objects which impact control flow
• [please confirm]: aux_data is deemed to be static, and therefore placing args which impact control flow in the tuple makes static_argnames / static_argnums unnecessary

[YouJiacheng]

Yes. aux_data is similar to the keys of dict and the order of tuple/list, i.e. intrinsically static. It is not only unnecessary to use static_arg* for pytree with proper aux_data, but also not correct -- unhashable children cannot be marked as static.

[ntenenz]

Understood. I was referring to the static_arg* being unnecessary for the fields listed in aux_data, which would otherwise have to be marked as static if the structure were flattened into a series of function arguments (as opposed to the argument container detailed above).

The connection I was missing is that jax leverages the intrinsically static nature of aux_data when tracing to follow a single code path. Thanks for clarifying!

[YouJiacheng]

Actually there is no magic.
JAX just replace all children with Tracer to record operations, and use fields in aux_data + tree-structure(treedef) + static attribute(shape, dtype) of children + static argument as key in cache of jit to avoid re-compilation.
The code execution is just pure python, JAX doesn't really leverage anything.