Are there any hidden side effects of having a list of jnp.ndarray as members to a struct.PyTreeNode

[sudo-michael]

The following code runs fine without any errors but I'm not sure why. I thought that all members of an struct.PyTreeNode must be immutable, but it appears that I can modify bar in-place. I also am surprised to not run into any jit errors.

I may not completely understand how PyTreeNode work but any explanation why this works / why this may not be a good idea would be greatly appreciated.

from flax import struct
import jax
import jax.numpy as jnp
class Foo(struct.PyTreeNode):
  bar: list[jnp.ndarray]
  @classmethod
  def create(cls, bar):
    return cls(bar)

@jax.jit
def my_sum(foo_state):
    s = 0
    for b in foo_state.bar:
        s += b.sum()
    return s
  
foo_state =Foo.create([jnp.array([1, 2, 3]), jnp.array([4, 5])])
print(my_sum(foo_state))
foo_state.bar.append(jnp.array([6]))
print(my_sum(foo_state))

[cgarciae]

While Foo is immutable, any mutable field can be updated given Python semantics. The main issue you will run into is that Foo is not hashable because list is not hashable. You might eventually get an error because of this.