How to deal with impure method with side-effect for registered Pytree class?

[HGangloff]

Hi all,

I have a follow-up MWE of the stateful computation tutorial where the Counter class is now in the pytree register, and I have a few question.

import jax
import jax.numpy as jnp

class Counter:
  """A simple counter."""

  def __init__(self, n=None):
    self.n = n if n is not None else 0

  def count(self) -> int:
    """Increments the counter and returns the new value."""
    self.n += 1
    return self.n

  def reset(self):
    """Resets the counter to zero."""
    self.n = 0

jax.tree_util.register_pytree_node(
    Counter,
    lambda self: ((self.n,), None), 
    lambda _, xs: Counter(xs[0]),
)

counter = Counter()

counter.reset()
# LOOP 1
@jax.jit
def loop_fun(counter):
    for _ in range(3):
        output = counter.count()
        jax.debug.print("{i}",i=output) # prints 1, 2, 3
loop_fun(counter)
# or scan equivalent
counter.reset()
# LOOP 2
_, res = jax.lax.scan(lambda c, _: (c, c.count()), init=counter, xs=jnp.arange(3)) 
print(res) # prints [1,2,3]

In loops 1 and 2, the loop iterations are taken into account by the jit call and we get the expected result. However, the count() function is impure because it modifies self which is not returned. And we know this is bad.

Are we getting the expected result by only chance ? Is returning self a good practice or is there a better way to make count() pure if we stick to the register_pytree_node-pattern ?

Thanks for your insights :)

[jakevdp]

Both cases are fine, becuase the functions you're using (loop_fun in the case of jit and the lambda in the case of scan) are non-side-effecting in themselves. The only side-effects that are problematic are those that happen across the boundaries of functions passed to transformations like jit or to higher-order primitives like scan.

You don't seem to be depending on the global state of counter being updated in either case, so your functions are effectively pure. Does that make sense?

[HGangloff]

Ok, this makes sense, thanks for the clear answer :)

[nicolasJouvin]

Hi, thanks for the question and the reply. I have an extra follow-up question elaborating on @HGangloff example but merely replacing the self.count attribute of Counter by a jnp.Array. When the bound method c.count() is jit-ed and evaluated before the scan, this leads to an UnexpectedTracerError. Why is that ? And why does this "jit-evaluate-then-scan" 3rd example behaves well when self.count is an int ?

import jax
import jax.numpy as jnp

class Counter:
  """A simple counter."""

  def __init__(self, n=None):
    self.n = n if n is not None else jnp.zeros((1,)) # use jnp.array

  def count(self) -> int:
    """Increments the counter and returns the new value."""
    self.n += 1
    return self.n

  def reset(self):
    """Resets the counter to zero."""
    self.n = jnp.zeros((1,)) # use jnp.array

jax.tree_util.register_pytree_node(
    Counter,
    lambda self: ((self.n,), None), 
    lambda _, xs: Counter(xs[0]),
)

counter = Counter()

counter.reset()
# LOOP 1
@jax.jit
def loop_fun(counter):
    for _ in range(3):
        output = counter.count()
        jax.debug.print("{i}",i=output) # prints 1, 2, 3
loop_fun(counter)
# or scan equivalent
counter.reset()
# LOOP 2
_, res = jax.lax.scan(lambda c, _: (c, c.count()), init=counter, xs=jnp.arange(3)) 
print(res) # prints [1,2,3]

# LOOP 3 - pre-jit bounded `counter.count()` method
counter.reset()
jitted_bounded_method = jax.jit(counter.count)
jitted_bounded_method() # call it once

_, res = jax.lax.scan(lambda c, _: (c, c.count()), init=counter, xs=jnp.arange(3)) # fails with UnexpectedTracerError

I tried to return self in the count` method like this (modifying the loop and scan accordingly), but this does not seem to help.

def count(self) -> int:
  """Increments the counter and returns the new value."""
  self.n += 1
  return self, self.n # returns self

For context : I am asking this to debug or rewrite a more elaborated code where a custom class (registered as PyTree) implements a method self.evaluate() modifying self.attribute within the call, where attribute is a PRNGkey array. The code fails with the same UnexpectedTracerError when a bounded method is evaluated through jax.values_and_grad(object.evaluate(*args)) during a scan.

[jakevdp]

This is working as expected. When you JIT-compile a bound method with a side-effect, you are implicitly attempting to persist that side-effect across the JIT boundary, which causes issues. The solution would be to not do that.

[nicolasJouvin]

Thank you for your answer. I did not get the problem of setting self.attr = ... inside a bound method when I first read Strategy 3 of the doc on JIT-compiling method. Is there any simple explanation or documentation as for why the following code actually works when using Python int instead of a JAX Array for Counter.n ? Thank you.

# @HGangloff MWE code here
# ...

# LOOP 3 - pre-jit bounded `counter.count()` method
counter.reset()
jitted_bounded_method = jax.jit(counter.count)
jitted_bounded_method() # call it once
_, res = jax.lax.scan(lambda c, _: (c, c.count()), init=counter, xs=jnp.arange(3)) # prints [1,2,3]

[jakevdp]

It "works" in the sense that it doesn't cache a traced value, because static integers are not traced.

It will not work in the sense that this side-effect only executes during trace-time, so you will not see the expected side-effect on subsequent calls if you run the function more than one time.