Downsides of preserving object identiy when flattening/unflattening

[liukidar]

Hello everyone,

in the past years I've been writing a lot of JAX code and I'm currently working on devolping a more object oriented interface to JAX (by basically having an abstraction layer that keep tracks of all the tensors through each jax transformation) and, among the many approaches I could think about, the following seems the easier. However, it goes against the core JAX idea that flattening/unflattening creates a copy of an object (i.e., a = unflatten(flatten(b)) => a is not b). So my question is, what are the problems that may arise with such approach?

In particular, you can look at the following (rather simplified) example:

import jax
from typing import Any, Tuple

class Param:
    def __init__(
        self,
        value
    ):
        self.value = value

    @staticmethod
    def flatten_parameter_with_keys(param: 'Param') -> Tuple[Any, 'Param']:
        return (("value", param.value),), param

    @staticmethod
    def unflatten_parameter(flatten_param: 'Param', data: Any) -> 'Param':
        param = flatten_param
        param.value = data[0]
        return param

jax.tree_util.register_pytree_with_keys(
    Param,
    flatten_with_keys=Param.flatten_parameter_with_keys,
    unflatten_func=Param.unflatten_parameter
)

@jax.jit
def increment(p):
    p.value = p.value + 1
    
    return p

a = Param(jax.numpy.array([1.0]))

increment(a)

print(a.value) # [2.]

Note how a.value is changed without having to reassign it.
I believe this approach, other than making everything stateful and object oriented, make also the overhead of jitting smaller as the flattening/unflattening operation of the function parameters and return values doesn't need to go through complex model, but just on the list of their parameters (as the model can simply be passed to the closure of the jitted function).

Of course if someone using this approach would want to actually create a copy of a Param using flattening/unflattening I could simply provide a function for it, so that's not a problem. I'm mostly interested in clashes with the internal functioning of jax and possible jitting optimizations that I may be missing.

Thank you for your time. I'm happy to have a deep discussion about this :)

[jakevdp]

The main downside here is that it would break some of the core assumptions of the APIs that use the pytree abstraction within JAX. Namely, the assumption that all relevant dynamic content is contained in the children, and that the aux_data is static, meaningfully hashable, and meaningfully compared using ==. This hash/equality property is used in e.g. the JIT cache to determine whether particular arguments to a JIT-compiled function require recompilation. As currently defined, instances of Param will never hash or compare as equal, so every time you use a new Param class as an argument to a jit-compiled function, it will trigger recompilation of that function.

The other thing I'd be worried about here would be leaked tracers. The approach is probably fine so long as pytree flattening/unflattening always happens in a paired fashion at a jit boundary, but that's not necessarily the case in JAX internals: there are many uses of pytrees where the resulting objects are passed around under the assumption that they only contain static information. I'm not immediately aware of any cases where this pattern would break things, but I would definitely not be surprised when it does.

[jakevdp]

recompilation is triggered only based on the value of "value"

value is an array specified as a dynamic argument (because it's part of children in the pytree lowering), so changing its value won't trigger recompilation: recompilation will only be triggered by changing its static compile-time attributes (shape, dtype, etc.)

indeed the instantiated object "param" is static (i.e. it doesn't change, only its content does, but this should be invisible to Jax)

What I mean is that param has non-static attributes. For example param.value is an array, and an array cannot be static. Including objects with non-static attributes in the pytree aux_data breaks the contract of pytree lowering, and I would not be surprised if this causes bugs. If you report those bugs, we'll respond that it's user error because you're breaking the contract of pytree lowering.

[liukidar]

value is an array specified as a dynamic argument (because it's part of children in the pytree lowering), so changing its value won't trigger recompilation: recompilation will only be triggered by changing its static compile-time attributes (shape, dtype, etc.)

Yes, this is exactly what I meant, basically the recompilation is triggered based on what "value" is and it is unaffected by Param. Given the definition of Param, this is exactly what happens since "value" is returned as its child and will trigger recompilation when necessary.

What I mean is that param has non-static attributes. For example param.value is an array, and an array cannot be static. Including objects with non-static attributes in the pytree aux_data breaks the contract of pytree lowering, and I would not be surprised if this causes bugs. If you report those bugs, we'll respond that it's user error because you're breaking the contract of pytree lowering.

I understand this bit. However, in my opinion, in the official documentation the static/dynamic splitting is not phrased as a strict requirement, but rather as a guideline to follow: "When defining an unflattening function, in general children should contain all the dynamic elements of the data structure (arrays, dynamic scalars, and pytrees), while aux_data should contain all the static elements that will be rolled into the treedef structure".

[jakevdp]

If it helps, you could replace “should” with “must” there. Regardless of the phrasing of the docs, what I’m saying is that the pattern you are using is dangerous: I wouldn’t recommend it, and things will likely go wrong.

[davisyoshida]

@liukidar The requirement to be static that Jake is warning about is relevant for a lot of things besides recompilation. For example, a lot of code relies on jax.tree_map making a copy of the input. Let's say you want to calculate the mean of each parameter while training:

import jax.numpy as jnp

def f(p):
  avgs = jax.tree_map(jnp.mean, p)
  return avgs

p = Param(jnp.arange(10))
f(p)
print(p.value) # D:

Any calculation you do like this will clobber the original values in the pytree. The same thing will happen if you just compute a grad and return it.

[liukidar]

The example you showed is exactly what I meant with "object oriented interface", but I agree it may be too counter-intuitive for the small benefits it offers. Thank you for your advice.