Internally, how do nnx wrappers interact with jax tracing to trace impure functions?

[hrbigelow]

Hi @cgarciae,

I wanted to get a basic understanding of how NNX allows an impure function to be traced by jax? Taking nnx.jit as an example:

@dataclasses.dataclass(eq=False)
class JitFn:
  f: tp.Callable[..., tp.Any]
  ...
  def __call__(self, *pure_args, **pure_kwargs):
    args, kwargs = extract.from_tree((pure_args, pure_kwargs), ctxtag='jit')

    out = self.f(*args, **kwargs)

    args_out, kwargs_out = extract.clear_non_graph_nodes((args, kwargs))
    pure_args_out, pure_kwargs_out, pure_out = extract.to_tree(
      (args_out, kwargs_out, out),
      ..
      ctxtag='jit',
      split_fn=_jit_split_fn,
    )

    return pure_args_out, pure_kwargs_out, pure_out

...
  jitted_fn = jax.jit(
    JitFn(fun, in_shardings, out_shardings, kwarg_shardings),
    in_shardings=jax_in_shardings,
    out_shardings=(jax_in_shardings, kwarg_shardings, jax_out_shardings),  # type: ignore
   ...
  )

The original non-pure function f is first wrapped in a function object, for example JitFn which accepts pure (constant) arguments and returns both f's output (in pure form) plus the possibly modified versions of each pure input. f is called within a graph context which monitors various things done during the call to f.

Mathematically this makes sense - you can express modifications to an argument as if it's a constant argument, and you return a modified copy of that input. But, ultimately, inside JitFn, f is still being called as-is, in unmodified form, complete with any statements which mutate variables:

out = self.f(*args, **kwargs)

What happens when jax.jit traces this part of the code? I mean, I understand conceptually how the outer function JitFn appears pure, but jax.jit actually traces individual statements, and so it requires the whole body of the function to not modify tensors.

Thanks again!

[cgarciae]

Hey @hrbigelow, this is a great question! Planning on documenting it in the near future as an advanced guided. Currently the best documentation for how this works internally is in the update_context docstrings.

flax/flax/nnx/graph.py

Lines 1060 to 1068 in 0679702

	idxmap
	(2) merge ─────────────────────────────► split (3)
	▲ │
	│ inside │
	│. . . . . . . . . . . . . . . . . . │ index_mapping
	│ outside │
	│ ▼
	(1) split──────────────────────────────► merge (4)
	refmap

Basically update_context keeps track of input/output references and their assigned integer index during the split called by to_tree internally, these indexes appear in the graphdef structure. The main trick is to have a mapping of which input reference indexes map to which output reference indexes (graphdef.index_mapping), this mapping is static it can be cached by jax.jit. Then the caller of JitFn can do a final extract.from_tree which will uses this information to figure out which output indexes on the graphdef correspond to which input references so merge merge can cache these references and use them instead of creating new ones but modify them exactly as specified by the output graphdef.

I'm omitting many details here but that is the core mechanism. In the end its just split + merge with some bookkeeping, but its surprisingly power and jax.jit friendly.

[hrbigelow]

Thanks @cgarciae,

Thanks - yes actually I did study the graph.py and extract.py code. it helped a bit, although some parts are still mysterious. One follow-up: is taking a merge of a split of some argument values restore those same values (aside from side effect changes to the graph context)?

That is:

round_trip = merge(split(original))
assert round_trip == original # equal values

And just to fix ideas maybe another way to write what you wrote, is this more or less the pattern? (pretend everything inside jax.jit { } has been transformed inline, I suppose).

def nnx_fn(*args):
  pure_args = split(args) # 1
  jax.jit {
    args = merge(pure_args) # 2
    out = fn(*args) # Does args have the same values as provided to nnx_fn?
    args_out = clear_non_graph_nodes(args)
    pure_out, pure_args_out = split(args_out, out) # 3
  }
  out = merge(pure_out) # 4
  return out

[cgarciae]

is taking a merge of a split of some argument values restore those same values (aside from side effect changes to the graph context)?

I think the proper way to describe merge(split(original)) is that you get an identical graph, equality for python objects is identity by default so you will get False even if they have the same values.

And just to fix ideas maybe another way to write what you wrote, is this more or less the pattern?

Yeah this pseudo code you showed is more or less what happens. However, maybe look at this this test

flax/tests/nnx/graph_utils_test.py

Line 537 in 0679702

def test_split_merge_update_context(self):

which manually uses update_context and graph.split/merge_context to recreate what nnx.jit is doing.

    class Foo(nnx.Module):
      def __init__(self):
        self.a = nnx.Param(1)
        self.b = 2

    m = Foo()
    ctxtag = 'test'

    with nnx.update_context(ctxtag):
      # roughly what to_tree is doing
      with nnx.graph.split_context(ctxtag) as ctx:
        # splitting m twice to simulate aliasing / reference sharing
        graphdef1, state1 = ctx.split(m)
        graphdef2, state2 = ctx.split(m)  # repeated

      @jax.jit
      def f(graphdef1, state1, graphdef2, state2):
        # roughly what from_tree is doing
        with nnx.graph.merge_context(ctxtag) as ctx:
          m1 = ctx.merge(graphdef1, state1)
          m2 = ctx.merge(graphdef2, state2)

        # swap a and b
        m1.a, m1.b = m1.b, m1.a

        with nnx.graph.split_context(ctxtag) as ctx:
          graphdef1, state1 = ctx.split(m1)
          graphdef2, state2 = ctx.split(m2)

        return graphdef1, state1, graphdef2, state2

      graphdef1, state1, graphdef2, state2 = f(
          graphdef1, state1, graphdef2, state2
      )

      with nnx.graph.merge_context(ctxtag) as ctx:
        m1_out = ctx.merge(graphdef1, state1)
        m2_out = ctx.merge(graphdef2, state2)
        
      assert m is m1_out
      assert m is m2_out

[hrbigelow]

I think the proper way to describe merge(split(original)) is that you get an identical graph, equality for python objects is identity by > default so you will get False even if they have the same values.

okay, quick follow up, when you say you get an identical graph, do you mean something like the container structure between original and round_trip are the same, and the values of the leaves are the same? i.e. something like you would get if you copy.deepcopy(original)?

[cgarciae]

Correct, you basically get a copy with the same structure and values but not the same references.

[hrbigelow]

I noticed this same 4-step box motif is happening in each of the nnx transformations pmap, vmap, jit. So, why is it necessary to have a merge and split being performed inside the given jax transformation? Is it something like, nnx is "observing" how the jax transformation interacts with the referential mutations that are happening inside the user-provided function?

My base-level understanding was that all jax transformations (vmap, pmap, jit, grad) assume the function is pure. So it is surprising that the nnx transformations actually all call the original user-provided impure function inside the jax transformation unmodified.

Is there some background (mathematical, or otherwise) that you think I might need in order to understand this basic "4-step box motif" (as I'll call it)? I'm really not grokking it. If it's a longer-term thing, I can wait for an advanced guide too.