Only consider Tracers with ShapedArray avals as instances of Array.

This fixes a TODO from #33420.

PiperOrigin-RevId: 835283505

Author: hawkinsp

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

@@ -2201,18 +2201,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;
       },