vmap crash with tree_flatten trigger

[jecampagne]

Hi,
here is a snippet that reproduce may concern using a third party that designed the equivalent of BaseObj:

class BaseObj():
    def __init__(self, *, gsparams=None, **params):
      self._params = params  # Dictionary containing all traced parameters
      self._gsparams = gsparams  # Non-traced static parameters

    @property
    def gsparams(self):
        """A `GSParams` object that sets various parameters relevant for speed/accuracy trade-offs."""
        return self._gsparams

    @property
    def params(self):
        """A Dictionary object containing all parameters of the internal represention of this object."""
        return self._params


    def tree_flatten(self):
        """This function flattens the GSObject into a list of children
        nodes that will be traced by JAX and auxiliary static data."""
        # Define the children nodes of the PyTree that need tracing
        children = (self.params,)
        # Define auxiliary static data that doesn’t need to be traced
        aux_data = {"gsparams": self.gsparams}
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        """Recreates an instance of the class from flatten representation"""
        return cls(**(children[0]), **aux_data)

@register_pytree_node_class
class T(BaseObj):
  def __init__(self, a, gsparams=None):
    super().__init__(a=a, gsparams=gsparams)

    # self._m = a*2     #<<<<<<<<<<<<<<  pb if decommented

  @property
  def a(self):
    return self.params["a"]
  
  def f(self,x):
    return jnp.sin(self.a*x)

If I instantiate a collection of T objs and apply a vmap

avals = jnp.pi * jnp.array([1.,2.])
ts = T(avals)
jax.vmap(lambda t: t.f(0.1))(ts)

Then if I need the self._m parameter to compute the result of f-function then I get a crash

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)

[<ipython-input-109-bb5a2ac90778>](https://localhost:8080/#) in <module>
----> 1 jax.vmap(lambda t: t.f(0.1))(ts)

[<ipython-input-93-9d50df48b013>](https://localhost:8080/#) in tree_unflatten(cls, aux_data, children)
     27     def tree_unflatten(cls, aux_data, children):
     28         """Recreates an instance of the class from flatten representation"""
---> 29         return cls(**(children[0]), **aux_data)

[<ipython-input-107-1cb043af2c87>](https://localhost:8080/#) in __init__(self, a, gsparams)
      4     super().__init__(a=a, gsparams=gsparams)
      5 
----> 6     self._m = a*2
      7 
      8   @property

TypeError: unsupported operand type(s) for *: 'object' and 'int'

Do you know how I can compute _m at inititialization and that I can use it in the T-functions later?
Thanks

[jakevdp]

Thanks for the question – this issue is covered in Custom PyTrees and Initialization. You need to allow for the fact that your pytrees may be instantiated with placeholder values. You can do that by explicitly checking for such values, or by avoiding calling __init__ within tree_unflatten and instead instantiating the object via object.__new__.

[soraros]

It would be nice if the placeholder objects have __repr__, do you think we can add that so the problem becomes easier to debug? If the object has a name, it's easier to search and find the relevant doc.

[jecampagne]

Nice @jakevdp that could be even nicer if you can write a possible solution in may case. Just to be sure that I do not make a mistake. Thanks in advance.

[soraros]

@jecampagne, check out @jakevdp's example of alternative for pytree initialization.

[jecampagne]

Not sur if I'm correct

class T(baseObj):
  def __init__(self, a):
    if not (type(a) is object or a is None or isinstance(a, MyTree)):
       _m = fcomplex_init(a)
    self.a = a
    self._m = _m ???

[soraros]

Since a is traced, and _m is a (pure) function of a, I think it's not important that the computation happens at __init__ time. So, other than what's suggested by @jakevdp, you could write

class T(BaseObj):
  def __init__(self, a, gsparams=None):
    super().__init__(a=a, gsparams=gsparams)

  @property
  def _m(self):
    return self.a * 2

  ...

  def f(self, x):
    return jnp.sin(self.a * x) + self._m  # <<< use it here

[soraros]

@jecampagne I think as long as it only depends on constant things, it will only run once in the jitted code (not entirely sure though). Let me write you an answer in a few days, being a bit busy lately.

[jecampagne]

@soraros just a up to your proposal to write an example...

[jecampagne]

@soraros hi! do you remember...

[soraros]

@jecampagne Sorry for being unresponsive. Will try to write something today!

[soraros]

@jecampagne I was wrong in thinking self._y evaluated only once. The easiest fix seems to be moving y's initialisation logic into the smart constructor init similar to the NumPy version.

Updated lines are marked with # <<< in the example.

Updated example

from functools import partial
from typing import Optional

import numpy as np

import numpy as np
import numpy.typing as npt

import jax
import jax.numpy as jnp
import jax.tree_util as tu
from jax import Array, jit, lax, vmap
from jax.typing import ArrayLike

rng = np.random.default_rng(0)
jax.config.update('jax_enable_x64', True)

class ExprNp:
  def __init__(self, name: str, steps: int, const: npt.NDArray,
               *, start: int = 0):
    self.name = name  # metadata
    self.steps = steps  # constant, static
    self.const = const  # constant, array
    # state, initialized in __init__, modified throughout
    self._start = start
    self._state = np.zeros_like(const, dtype=np.float_)
    # value, depends on constant, initialized in __init__, never modified again
    self._y = self._init_y(const)
    # state, depends on state, initialized in __init__, modified throughout
    self._sum = np.zeros(steps)

  @property
  def state(self):
    return self._state, self._sum

  @staticmethod
  def _init_y(c: npt.NDArray):
    return np.sin(c)

  def _step(self, i):
    self._state += self._y
    self._sum[i] += 1

  def main(self, steps: Optional[int] = None, *, start: Optional[int] = None):
    steps = steps or self.steps
    start = start or self._start
    end = start + steps
    for i in range(start, end):
      self._step(i)
    self._start = end
    return self

@tu.register_pytree_node_class
class ExprJax:
  def __init__(self, name: str, steps: int, const: Array, start: Array,
               _state: Array, _sum: Array, _y: Array):
    self.name = name
    self.steps = steps
    self.const = const
    self._start = start
    self._state = _state
    self._sum = _sum
    self._y = _y  # <<< add a field to store `y`

  @classmethod
  def init(cls, name: str, steps: int, const: ArrayLike, *, start: int = 0):
    _state = jnp.zeros_like(const, dtype=np.float_)
    _start = jnp.full(_state.shape[:-1], start)
    _sum = jnp.zeros(_state.shape[:-1] + (steps,))
    _y = cls._y_init(const)  # <<< initilize `y` in the smart constructor
    return cls(name, steps, jnp.array(const), _start, _state=_state, _sum=_sum, _y=_y)

  @property
  def state(self):
    return self._state, self._sum

  def tree_flatten(self):
    # <<< also passing `y` to `tree_flatten` to make sure it's not re-initialized
    return ((self.const, self._start, self._state, self._sum, self._y), 
            (self.name, self.steps))

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    return cls(*aux_data, *children)

  @staticmethod
  def _y_init(const: ArrayLike):
    jax.debug.print('y')  # <<< print to verify it's only called once
    return jnp.sin(const)

  def _step(self, i: Array):
    self._state += self._y
    self._sum = self._sum.at[i].add(1)

  @partial(jit, static_argnums=(1,))
  def main(self, steps: Optional[int] = None, *, start: Array = jnp.array(0)):
    assert self._state.ndim == 1
    steps = steps or self.steps
    start = start or self._start
    def body(i: int, self: ExprJax):
      self._step(start + i)
      return self
    self = lax.fori_loop(0, steps, body, self)
    self._start = start + steps
    return self

data = np.arange(10).reshape(2, -1)

res_np = [ExprNp('sim_np', 10, d).main(3).state for d in data]
res_jax = [ExprJax.init('sim_jax', 10, d).main(3).state for d in data]

sim_jax = ExprJax.init('sim_jax', 10, data)
res_mapped = jit(vmap(partial(ExprJax.main, steps=3)))(sim_jax)

def tree_allclose(a, b):
  return all(tu.tree_leaves(tu.tree_map(np.allclose, a, b)))

assert tree_allclose(res_jax, res_np)
assert tree_allclose(res_mapped.state, tu.tree_map(lambda *a: np.stack(a), *res_np))

run3 = jit(vmap(partial(ExprJax.main, steps=3)), donate_argnums=(0,))
run6 = jit(vmap(partial(ExprJax.main, steps=6)), donate_argnums=(0,))

sim_jax33 = run3(run3(ExprJax.init('sim_jax', 10, data)))
sim_jax6 = run6(ExprJax.init('sim_jax', 10, data))

assert tree_allclose(sim_jax33.state, sim_jax6.state)