Debugging ValueError: JaxprInputEffect Read<5> is invalid.

[pfrommerd]

When using nnx with the new array_ref API, it appears there are some bugs with some jax transformations as I get

ValueError: JaxprInputEffect Read<5> is invalid.

followed by an extremely large jaxpr. It doesn't happen with all nnx models that I use and I am struggling to produce a reproducible minimal example I can post here--I don't have any issues when using the old-style nnx API.

I also sometimes get a similar-looking error that suggests not enough inputs to a jaxpr for a particular Read<> effect.
I suspect this is an internal bug with some jax transformation that I am doing and how it interacts with array_refs (i.e. something similar to #16370 potentially, although that specific issue has been marked as closed)

Does anyone have suggestions as to how I can debug this in a manageable manner? Is there some flag I can enable for super-stringent jaxpr validation that will catch which transformation is causing the issue?

[pfrommerd]

Okay I think I have managed to create a reproducible example of code that breaks.

It appears that are still issues with jax.lax.scan, because the following code:

import jax
import jax.numpy as jnp


@jax.jit
def bar(x, w):
    def scan_fn(x, _):
        c = jnp.array([])
        x = w[...] @ x
        x = jnp.concatenate([x, c], axis=-1)
        return x, None

    x, _ = jax.lax.scan(scan_fn, x, None, length=10)
    return x


@jax.jit
def foo(w):
    return bar(
        jnp.zeros((1,)),
        w,
    )


foo(jax.array_ref(jnp.eye(1)))

Gives the error:

ValueError: JaxprInputEffect Read<2> is invalid.

Curiously, using jnp.zeros((0,)) instead of jnp.array([]) fixes the issue.
Perhaps this has something to do with the new handling of array literals?

I will open a bug report.

[pfrommerd]

I have created #32399

[mattjj]

Thanks so much for flagging this!

[pfrommerd]

Thanks for patching this so quickly! I can confirm that I no longer have issues with the example I posted above, but am still running into issues with the full model (the same error and the jaxpr involves consts so I assume a similar bug elsewhere). I will track down the issue and open a new bug when I can reproduce with a simple example.

[mattjj]

Sorry the bugs persist! Yes probably a similar issue. If you set jax.config.update(‘jax_enable_checks’, True) and jax.config.update(‘jax_traceback_filterint’, ‘off’) and share the traceback, that might be enough to figure it out. The minimal repros are extremely helpful though.

[pfrommerd]

Thanks for the tip! The issue now seems to occur due to jax.lax.cond and effects not properly handling constvars. Here is another minimal repro that has issues on nightly:

import jax
import jax.numpy as jnp

@jax.jit
def bar(w):
    x = jnp.zeros((1,)) + jnp.array([0,])
    x = jax.lax.cond(x[0] < 1, lambda x: x + w[...], lambda x: x - 1, x)
    return x


@jax.jit
def foo(w):
    return bar(w)


foo(jax.array_ref(jnp.ones((1,))))

[mattjj]

Nice job tracking it down!

Could you open a new issue so I don’t forget?

[pfrommerd]

Opened #32446