Problem in `_Quaxify.__call__`

[nstarman]

_Quaxify.__call__ is converting things to tracers when they shouldn't be tracers. For example, here's a stable MWE

import jax.numpy as jnp
import quax

xbool = jnp.array([ True, False,  True], dtype=bool)
x1 = jnp.array([1., 2., 3.], dtype=float)

compress = quax.quaxify(jnp.compress)
compress(xbool, x1) 

Here's the traceback. I think this is saying that some stuff are becoming tracer, and that's triggering the error.

_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../quax/quax/_core.py:311: in __call__
    out = fn(*args, **kwargs)
        args       = (Array([ True, False,  True], dtype=bool), Array([1., 2., 3.], dtype=float32))
        dynamic    = (<function compress at 0x1084b2340>, (Array([ True, False,  True], dtype=bool), Array([1., 2., 3.], dtype=float32)), {})
        fn         = <function compress at 0x1084b2340>
        kwargs     = {}
        parent_trace = EvalTrace
        self       = _Quaxify(fn=<function compress>, filter_spec=True, dynamic=False)
        static     = (None, (None, None), {})
        tag        = <jax._src.core.TraceTag object at 0x110d3be90>
        trace      = _QuaxTrace
.venv/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py:9595: in compress
    if reductions.any(extra):
        a          = Array([1., 2., 3.], dtype=float32)
        arr        = Traced<ShapedArray(float32[3])>with<_QuaxTrace> with
  value = _DenseArrayValue(array=f32[3])
        axis       = 0
        condition  = Array([ True, False,  True], dtype=bool)
        condition_arr = Array([ True, False,  True], dtype=bool)
        extra      = Traced<ShapedArray(bool[0])>with<_QuaxTrace> with
  value = _DenseArrayValue(array=bool[0])
        fill_value = Traced<ShapedArray(int32[], weak_type=True)>with<_QuaxTrace> with
  value = _DenseArrayValue(array=weak_i32[])
        out        = None
        size       = None
.venv/lib/python3.11/site-packages/jax/_src/core.py:842: in __bool__
    return self.aval._bool(self)
        self       = Traced<ShapedArray(bool[])>with<_QuaxTrace> with
  value = _DenseArrayValue(array=bool[])
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = ShapedArray(bool[]), arg = Traced<ShapedArray(bool[])>with<_QuaxTrace> with
  value = _DenseArrayValue(array=bool[])

    def error(self, arg):
>     raise TracerBoolConversionError(arg)
E     jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
E     See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

arg        = Traced<ShapedArray(bool[])>with<_QuaxTrace> with
  value = _DenseArrayValue(array=bool[])
self       = ShapedArray(bool[])

[nstarman]

I think this is the more important thing to fix, over #57 which I can hack around.
I haven't figured out a way around this problem.

[patrick-kidger]

I think this is a property of jnp.compress, not Quax:

import jax
import jax.numpy as jnp

xbool = jnp.array([ True, False,  True], dtype=bool)
x1 = jnp.array([1., 2., 3.], dtype=float)

@jax.jit
def compress(condition, a):
    return jnp.compress(condition, a)

compress(xbool, x1)

the output size is dependent on the value of condition. Passing the size argument has the above work for both JAX and Quax: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.compress.html

[nstarman]

But quaxify(jnp.compress()) doesn't jit the function so where jnp.compress works without the size argument, shouldn't quaxify(jnp.compress) ?

[patrick-kidger]

Ah right. So I'm guessing we should add the same thing that e.g. JVPTracer.__bool__ has:

    if is_concrete(self): return bool(self.to_concrete_value())  # pytype: disable=wrong-arg-types

[nstarman]

I'm looking at _QuaxTracer and it directly inherits from core.Tracer which I think gets the same __bool__ as JVPTracer through the common Array base class. So _QuaxTracer has this already.

@patrick-kidger Were you thinking of adding this to _Quaxify or _DenseArrayValue?

Signature: quax._core._QuaxTracer.__bool__(self)
Docstring: <no docstring>
Source:   
  def __bool__(self):
    if is_concrete(self): return bool(self.to_concrete_value())  # pytype: disable=wrong-arg-types
    check_bool_conversion(self)
    return self.aval._bool(self)
File:      [~/local/quax/.venv/.../python3.10/site-packages/jax/_src/core.py)
Type:      function

[nstarman]

In GalacticDynamics/quaxed#136 I identify the full set of numpy functions impacted by this problem. xfail_quax58 is used to filter the tests.

xfail_quax58 = pytest.mark.xfail(
     reason="https://github.com/patrick-kidger/quax/issues/58"
 )

[nstarman]

I'm less well versed in how JAX turns arrays into tracers, so the point at which I'm getting confused is how:

.venv/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py:9595: in compress
    if reductions.any(extra):
        extra      = Traced<ShapedArray(bool[0])>with<_QuaxTrace> with
  value = _DenseArrayValue(array=bool[0])

But if I use --pdb I find that extra is a Tracer but extra = condition_arr[arr.shape[0]:] from the line before is concrete.
condition_arr is concrete (even in the stacktrace) and arr.shape has to be concrete. So how is extra turning in to a Tracer, causing it to fail reductions.any?

I haven't figured out what in quax is causing this spurious transformation.

[patrick-kidger]

I'm looking at _QuaxTracer and it directly inherits from core.Tracer which I think gets the same bool as JVPTracer through the common Array base class. So _QuaxTracer has this already.

Ah, you're right, I was mistaken. So I guess the question is why we're not passing the is_concrete check then.

Which I can see is what you're starting to think about in the follow-up message.

---

FWIW it's not clear to me that this is too bad of a problem. If under jit then these would fail anyway, and JAX code is(/should) nearly always be ran under jit. Do you have use-cases for unjitted code?

[nstarman]

FWIW it's not clear to me that this is too bad of a problem

I broadly agree, however it does impact quaxed which is meant to be a drop-in replacement for jax functions.

[nstarman]

I think the source of the issue is the jax.jit in
https://github.com/patrick-kidger/quax/blob/b72049dd8afdc951abb4ac3b2198cfc638f5bbbe/quax/_core.py#L541C1-L552C70

Can the jit be avoided when constructing a call to pjit_p?

E.g. the following doesn't correctly unflatten the arguments, so return types are wrong, but it does call out to pjit_p

leaves, treedef = jtu.tree_flatten(args)  # remove all Values
return jax._src.pjit.pjit_p.bind(*leaves, jaxpr=jaxpr, inline=inline, **kwargs)

[patrick-kidger]

Ah, interesting! Yeah, I'm definitely willing to believe that the nested jax.jit is the root cause of some funny business here.

[nstarman]

The current solution is quite clever in that it moves the quax stuff to inside the jit. Then the Jax arguments can be passed through and the quaxed type returned, without any recursion problems. Can we do something similar with the pjit_p.bind and a jaxpr? Could we take the jaxpr, turn it into a function. Wrap that function in something that unflattens the arguments into their quaxed values, and turn that wrapper function into a new jaxpr which can be passed to pjit_p.bind?

Asking ChatGPT for some help with the above text as input and the current function as context... [then modified]

def symbolic_like(x):
    if isinstance(x, jax.Array):
        return ShapeDtypeStruct(x.shape, x.dtype)
    elif isinstance(x, (bool, int, float)):  # Python scalars
        # Wrap with jnp to infer dtype
        arr = jnp.array(x)
        return ShapeDtypeStruct((), arr.dtype)
    else:
        raise TypeError(f"Cannot generate symbolic placeholder for {type(x)}")


@register(jax._src.pjit.pjit_p)  # pyright: ignore
def pjit_quax(
    *args: ArrayLike | ArrayValue, jaxpr: Any, inline: bool, **kwargs: Any
) -> Any:
    fun = quaxify(jexc.jaxpr_as_fun(jaxpr))
    if inline:
        del jaxpr, kwargs
        return fun(*args)

    # Flatten input arguments
    leaves, treedef = jtu.tree_flatten(args)

    # Define a function that reconstitutes the original tree
    def wrapped_fun(flat_args):
        return fun(*jtu.tree_unflatten(treedef, flat_args))

    # Generate example inputs: ShapeDtypeStructs for symbolic jaxpr tracing
    example_inputs = jtu.tree_map(symbolic_like, leaves)

    # Get a jaxpr that takes flattened inputs
    wrapped_jaxpr = jax.make_jaxpr(wrapped_fun)(example_inputs)

    return quaxify(jax._src.pjit.pjit_p.bind)(
        *leaves, jaxpr=wrapped_jaxpr, inline=True, **kwargs
    )

That doesn't seem to work. The output type is always a JAX array, not the intended array-ish object.