Spec undefined behaviour: cloning semantics for jit-related functions

[cgarciae]

Hey JAX team!

@shoyer was kind enough to point out that I am depending on undefined behaviour on this twitter thread when updating incoming Pytrees in place. Example:

jax.jit
def train_step(model, x, y, optimizer):
    ...
    params = optimizer.update(grads, params)
    ...
    return model, loss, optimizer

The following is happening:

1. optimizer is a Pytree so it passes through jit.
2. update is mutating optimizer in place
3. The previous currently works because jit effectively clones all pytrees every time when flattening / unflattening.
4. optimizer is output by train_step, thus propagating the mutated state.

The problem I was not aware of is that 3 is undefined behaviour as @shoyer tells me, current version of JAX effectively clone the Pytrees every time so the example works but is not guaranteed. I want to point out the following:

• I don't know the internal implications of this, performance or otherwise.
• Cloning seems to be in general a safer option for users, might help protect them by being defensive.

Is it worth creating a spec for this? Too early to know?

Thanks!

[cgarciae]

BTW: a discussion around this would immensely help to stir Treex's API as state management is a core issue.

https://github.com/cgarciae/treex

[jekbradbury]

I think this specific code is okay. In general performing in-place updates of Python objects is fine as long as you:

• always pass such objects as explicit arguments and return values of every function that you wrap with a JAX transformation like jit, vmap, scan,named_call, etc. Failing to do this violates the "contract" of JAX transformations (they need the functions they transform to be functionally pure), and it's this requirement that gives rise to the wrapped transformation APIs in libraries like haiku, flax, and objax (which maintain state in mutable Python objects, but then lift these objects to be explicit arguments and return values of transformed functions).
• never rely on the object identity of mutable Python objects inside functions you'd like to wrap with JAX transformations (e.g. don't write things like if a is b or create dictionaries with object-identity keys). The JAX pytree infrastructure will not necessarily preserve these identity relationships.

These two things aren't directly related, they're just both constraints relating to in-place object manipulation.

[soraros]

I think this answer would be a nice addition to The Sharp Bits.
The section Pure functions does mention using stateful objects internally, but not so detailed.

[jakevdp]

I've opened #7759 to track adding these docs; assigned to James for now.

[cgarciae]

I see. Thanks @jekbradbury!

I've been intuitively counting on the rules you describe. Are they established somewhere in the docs? If not maybe adding a brief section about "Object mutation inside JIT operations" with this information would be valuable.