Fix _QuaxTracer concrete bool conversion for plain JAX arrays

[Copilot]

_Quaxify.__call__ wraps all dynamic JAX arrays as _QuaxTracer objects inside a _QuaxTrace context. When the wrapped function internally requires Python-level boolean evaluation of a concrete array (e.g. jnp.compress doing if reductions.any(extra):), it fails with TracerBoolConversionError because _QuaxTracer never overrode to_concrete_value(), so is_concrete() always returned False.

# Previously raised TracerBoolConversionError
compress = quax.quaxify(jnp.compress)
compress(jnp.array([True, False, True]), jnp.array([1., 2., 3.]))
# Now works correctly → Array([1., 3.], dtype=float32)

Changes

• src/quax/_core.py — Override to_concrete_value() on _QuaxTracer: when the wrapped value is a _DenseArrayValue, delegate to core.to_concrete_value(self.value.array). This allows __bool__ (and __array__) to resolve concretely for plain JAX arrays passing through quaxify, while correctly returning None under JIT.

• tests/unit/test_core.py — Add test_concrete_bool_conversion reproducing the issue MWE.

• tests/unit/test_numpy/test_jax_array.py — Remove xfail_quax58 from the 15 functions that now pass without a size argument: argwhere, bincount, choose, compress, flatnonzero, nonzero, ravel_multi_index, roots, trim_zeros, union1d, unique* (5 variants). Keep xfail_quax58 on extract, intersect1d, setdiff1d, setxor1d (still fail with TracerArrayConversionError).

• tests/unit/test_numpy/test_myarray.py — Remove xfail_quax58 from choose (only case that passes; all others still fail because the condition is a custom Value subclass, not a _DenseArrayValue).

[nstarman]

Excellent. A simple solution. Most tests are passing. Not all, so let's not close #58 quite yet.

[Copilot]

Pull request overview

This PR fixes a TracerBoolConversionError when quaxify-wrapped functions attempt Python boolean evaluation on concrete (eager) JAX arrays by making _QuaxTracer able to produce a concrete value for _DenseArrayValue-backed arrays.

Changes:

• Override _QuaxTracer.to_concrete_value() to delegate concretization for _DenseArrayValue payloads.
• Add a regression test covering concrete boolean conversion via jnp.compress.
• Remove xfail_quax58 from multiple NumPy-on-JAX tests that now pass; keep xfails for cases still failing with array conversion errors.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
src/quax/_core.py	Adds _QuaxTracer.to_concrete_value() override to enable concretization for plain JAX arrays.
tests/unit/test_core.py	Adds regression test reproducing the concrete-bool conversion failure (issue #58).
tests/unit/test_numpy/test_jax_array.py	Un-xfails a set of jnp.* functions now passing without static size.
tests/unit/test_numpy/test_myarray.py	Un-xfails choose for MyArray where it now passes.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[nstarman]

@meeseeksdev help

[nstarman]

@meeseeksdev hello

[lumberbot-app]

Helloooo @nstarman, I'm Mr. Meeseeks! Look at me!