Any reason for jax.lax.cond not playing nice with jacobian

[samskiter]

I'm currently trying to reproduce a simpler version of a bug I'm facing in my code - essentially a jax.lax.cond is killing my gradients in my jacobian. Before I go mad, are there any gotchas I might be missing?

I have a class that implements a method something like this:

def get_error(self) -> float:

    return jax.lax.cond(
        self.some_class_value > 1.0,
        self._calculate_error,
        self._calculate_other_error,
    )

This is used in some higher error function, which I then take a jacfwd of.

I noticed I was seeing zero gradients in parts of my jacobian in places where I saw them on an old reference implementation. I also noticed that if I removed the conditional, those gradients came back and matched the reference implementation.

The following worked (In my debugging case I knew the conditional evaluated false):

def get_error(self) -> float:
    self._calculate_other_error()

The following did not work:

def get_error(self) -> float:

    return jax.lax.cond(
        False,
        self._calculate_error,
        self._calculate_other_error,
    )

Nor did the following:

def get_error(self) -> float:

    return jax.lax.cond(
        True,
        self._calculate_other_error,
        self._calculate_other_error,
    )

It seems that merely the presence of the cond is throwing something off in the jacobian

As I mentioned, I'm still trying to reproduce this outside of my code base (where things seem to be behaving a little better), but before I go mad, are there any gotchas I ought to be looking out for ?

Thanks!

[jakevdp]

What do the false branches return? This might be some manifestation of this issue JAX FAQ: gradients contain NaN when using where. The gradient of a cond is basically computed as 1 * grad(true_branch)(x) + 0 * grad(false_branch(x)), so if false_branch returns something like NaN or inf, it can cause issues with the overall gradient.

In any case, if you're able to put together a minimal reproduction, it would be easier to help.

[samskiter]

Thanks for the reply - still trying to get this minimal repro together - I read the FAQ article you linked to but thats for reverse mode differentiation - I'm using jacfwd (infact I've tried with both)

Anyway, the branch just returns a number: 0.09181296626905545 to be precise, so it's not NaN or inf.

[jakevdp]

Another guess: does your class do any in-place mutation (i.e. methods that have side-effects of the form self.attr = ...). If your functions are impure, I could imagine that causing this kind of issue in autodiff.

[samskiter]

Definitely have lots of impure stuff but I've registered pytrees. Flatten/unflatten not being called though :/

…

On Tue, 29 Aug 2023, 22:10 Jake Vanderplas, ***@***.***> wrote: Another guess: does your class do any in-place mutation (i.e. methods that have side-effects of the form self.attr = ...). If your functions are impure, I could imagine that causing this kind of issue in autodiff. — Reply to this email directly, view it on GitHub <#17341 (reply in thread)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJWOW23E2OCYYH455GV3TTXXZLE5ANCNFSM6AAAAAA4CD76NU> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[jakevdp]

Registering pytrees is not enough to make code work correctly in the presence of mutation. JAX transforms still rely on functions being pure. I strongly suspect this is the root of your issue.

[samskiter]

Interesting. My code looks roughly like:

class Thing:
    foo: float
    def get_error() -> float:
        return f(self.foo)

class Optimizer:

    def __init__(self, things: [Thing]):
        self.things = things
        self.variables = [0.1] * len(things)

    def get_errors(things, variables):
        errors = jnp.empty(len(self.things)
        for idx, thing in enumerate(things):
            thing.foo = variables[idx] # <--- would this qualify as the side-effect ?
        for idx, thing in enumerate(things):
            errors = errors.at[idx].set[things.get_error()]
        return errors

    def jacobian(things, variables):
        return jaxfwd(Optimizer.optimise)(things, variables)

    def optimize_step(self):
        errors = Optimizer.get_errors(self.things, self.variables)
        jacobian = Optimizer.jacobian)(self.things, self.variables)
        inv_jacobian = jnp.pinv(jacobian)
        correction = jnp.matmul(inv_jacobian, errors)
        self.variables -= correction
        return errors.sum()

    def optimize(self)
        while self.optimize_step() < 1e-7
            print("optimize!")

I've tried to make the above 'pure'functions within the Optimizer class, but it is still storing variables on the Thing classes prior to calling a function on them.... I tried something like this setup in a minimal repro but didn't see the issue....

[jakevdp]

# <--- would this qualify as the side-effect ?

Yes, definitely.

If you're updating attributes of self but not returning self from your function, then your function is impure and I would not expect JAX transformations to work with your code.

[samskiter]

@jakevdp yea, that's it - just nailed it down to a repro (below)! I wasn't expecting this to cause an issue after I had registered pytrees - my broad understanding was that allows my classes to be treated as if they were pure...

from __future__ import annotations
import jax
import jax.numpy as jnp
from jax.tree_util import register_pytree_node_class


@register_pytree_node_class  # <--- Even though Thing is registered as a pytree node, the jacobian function still fails
class Thing:
    foo: float

    def __init__(self, foo: float) -> None:
        self.foo = foo

    def tree_flatten(self) -> tuple[tuple[float], dict[str, float]]:
        return (self.foo,), {}

    @classmethod
    def tree_unflatten(
        cls, aux_data: dict[str, float], children: tuple[float]
    ) -> Thing:
        del aux_data
        return cls(children[0])

    def f(self, x: float) -> float:
        return x**2

    def my_f(self) -> float:
        return self.f(self.foo)

    def get_error(self) -> float:
        # return self.f(self.foo) # <--- this works
        return jax.lax.cond(  # <--- this fails
            True,
            self.my_f,
            self.my_f,
        )


class Optimizer:
    def __init__(self, things: [Thing]):
        self.things = things
        self.variables = jnp.array([0.1] * len(things))

    @staticmethod
    def get_errors(variables: jax.Array, things: list[Thing]) -> jax.Array:
        errors = jnp.empty(shape=(len(things)))
        for idx, thing in enumerate(things):
            thing.foo = variables[idx]  # <--- would this qualify as the side-effect ?

        print(f"things: {things}")
        for idx, thing in enumerate(things):
            error = thing.get_error()
            errors = errors.at[idx].set(error)
        return errors

    @staticmethod
    def jacobian(variables: jax.Array, things: [Thing]) -> jax.Array:
        return jax.jacfwd(Optimizer.get_errors, argnums=(0,))(variables, things)

    def optimize_step(self) -> float:
        print(f"variables: {self.variables}")
        print(f"things: {self.things}")
        errors = Optimizer.get_errors(self.variables, self.things)
        jacobian = Optimizer.jacobian(self.variables, self.things)
        print("jacobian", jacobian)
        exit(0)
        inv_jacobian = jnp.linalg.pinv(jacobian)
        correction = jnp.matmul(inv_jacobian, errors)
        self.variables -= correction
        return errors.sum()

    def optimize(self):
        while self.optimize_step() < 1e-7:
            print("optimize!")


optimizer = Optimizer([Thing(foo=0.1), Thing(foo=10.0)])
optimizer.optimize()

[samskiter]

Hrmm - yea re-reading the Jax documentation this is clear - often the docs are centered around @jit, but is then lumped in with grad and vmap, so I read the following sentence:

This should work correctly with JIT and other transforms so long as you never mutate your object. Mutations of objects used as hash keys lead to several subtle problems, which is why for example mutable Python containers (e.g. dict, list) don’t define hash, while their immutable counterparts (e.g. tuple) do.
<
If your class relies on in-place mutations (such as setting `self.attr = ... within its methods), then your object is not really “static” and marking it as such may lead to problems. Fortunately, there’s another option for this case.

Strategy 3: Making CustomClass a PyTree
The most flexible approach to correctly JIT-compiling a class method is to register the type as a custom PyTree object; see Extending pytrees. This lets you specify exactly which components of the class should be treated as static and which should be treated as dynamic. Here’s how it might look:

and this:

Note that grads is a pytree with the same structure as params.
jax.grad is one of the many JAX functions that has
built-in support for pytrees.

(from different sections of the manual)

I got confused and thought pytrees would solve everything!

[samskiter]

#1967 similarly

[jakevdp]

Yeah, I see where that could be confusing. But there's no getting around JAX's requirement that functions be pure.

[samskiter]

Just an update - while what I'm doing may not be totally 'kosher' as per @jakevdp 's replies above, I think it's worth showing a little bit of the techniques I used to migrate an OO codebase to Jax. In the ends I created superclass to inherit and follow a pattern to allow my OO objects to be correctly flattened. This works with transient dependencies and is closely related to a previous proposal for pydags, here: #7919

Some notes on the approach:

1. I had to use the register_pytree_node_class approach on all my types. I used a common superclass to make this easier: Posted here: PyDags proposal: a mechanism to preserve relative identities for Pytrees #7919 (comment)
2. The above implement a flatten and unflatten function and place the 'uuid's of children into the auxilliary data.
3. using grad(f) didn't work but grad(jit(f)) did. Seems the JIT function is essential for forcing the flatten to be applied
4. I also used VMAP successfully, but I had to make sure I was doing vmap(f(x)) and not cmap(f(x,y)) - i.e. that there was only one argument - and that was a registered pytree that contained all the objects associated together
5. As we've mentioned in the other comments, mutations to objects will not persist outside of a jitted function - therefore you must return any values you want out of them. This suited me fine - I have objects with fixed properties and want to mutate some to produce an error function and perform gradient descent.

I'm actually overall pretty happy with what I've ended up with - while I wouldn't start from this place - adapting a code base to here has left me with something that's still quite neat without any really nasty hacks. The flatten/unflatten logic is still contained within each subclass and separates concerns nicely. Copilot is now able to write all my flatten/unflatten and getrefs/setrefs functions.

I would be interested in seeing something like this brought to Jax itself to make migration to Jax a little easier...