Only consider Tracers with ShapedArray avals as instances of Array.

This fixes a TODO from #33420.

PiperOrigin-RevId: 861851484

Author: hawkinsp

CHANGELOG.md

@@ -20,6 +20,9 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
 * Bug fixes:
 * Deprecations:
 * Changes:
+  * JAX tracers that are not of `Array` type (e.g., of `Ref` type) will no
+    longer report themselves to be instances of `Array`.
+
 
 ## JAX 0.9.0 (January 20, 2026)
 

jax/_src/pallas/mosaic/pallas_call_registration.py

@@ -48,6 +48,9 @@ def _maybe_cast_to_int(x: jax.Array | jax_core.AbstractValue):
   """
   assert isinstance(
       x, (jax.Array, jax_core.ShapedArray, state_types.AbstractLinVal)
+  ) or (
+      isinstance(x, jax_core.Tracer)
+      and isinstance(x.aval, state_types.AbstractLinVal)
   ), type(x)
   if isinstance(x, jax.Array):
     if dtypes.issubdtype(x.dtype, jax.numpy.bool_):

jaxlib/py_array.cc

@@ -2221,18 +2221,11 @@ absl::Status PyArray::Register(nb::module_& m) {
         if (tracer_class.ptr() && self.ptr() == base_type.ptr() &&
             PyObject_TypeCheck(x.ptr(), reinterpret_cast<PyTypeObject*>(
                                             tracer_class.ptr())) != 0) {
-          // TODO(phawkins): we would like to change this to use the logic below
-          // but it is a somewhat breaking change. Let us defer it to a future
-          // PR.
-          return true;
-          // auto is_traced_array_fn =
-          //     nb::getattr(x, "_is_traced_array", nb::none());
-          // if (!is_traced_array_fn.is_none()) {
-          //   try {
-          //     return nb::cast<bool>(is_traced_array_fn());
-          //   } catch (...) {
-          //   }
-          // }
+          auto is_traced_array_fn =
+              nb::getattr(x, "_is_traced_array", nb::none());
+          if (!is_traced_array_fn.is_none()) {
+            return nb::cast<bool>(is_traced_array_fn());
+          }
         }
         return false;
       },