fix: filter also for new (v0.7.0) LinearizeTracer type (#3586)

filter also for new (v0.7.0) LinearizeTracer type

Author: pfackeldey

src/awkward/_kernels.py

@@ -6,6 +6,8 @@
 from abc import abstractmethod
 from typing import Any, Callable
 
+from packaging.version import parse as parse_version
+
 import awkward as ak
 from awkward._nplikes.cupy import Cupy
 from awkward._nplikes.jax import Jax
@@ -102,16 +104,20 @@ def __init__(self, impl: Callable[..., Any], key: KernelKeyType):
 
         self._jax = Jax.instance()
 
+        jax_module = ak.jax.import_jax()
+        self._ad_tracer_types = (jax_module._src.interpreters.ad.JVPTracer,)
+        if parse_version(jax_module.__version__) >= parse_version("0.7.0"):
+            self._ad_tracer_types += (jax_module._src.interpreters.ad.LinearizeTracer,)
+
     def _cast(self, x, t):
         if issubclass(t, ctypes._Pointer):
             # Do we have a JAX-owned array?
             if self._jax.is_own_array(x):
                 if self._jax.is_tracer_type(type(x)):
-                    jax_module = ak.jax.import_jax()
                     # general message for any invalid JAX input type
                     msg = f"Encountered {x} as an (invalid) input to the '{self._key[0]}' Awkward C++ kernel."
                     # message specification for autodiff (i.e. when encountering a JVPTracer)
-                    if isinstance(x, jax_module._src.interpreters.ad.JVPTracer):
+                    if isinstance(x, self._ad_tracer_types):
                         msg += " This kernel is not differentiable by the JAX backend."
                     raise ValueError(msg)
                 assert self._jax.is_c_contiguous(x), "kernel expects contiguous array"