Repository: scikit-hep/awkward

Issue:
Jax tracer TypeError in multiplication
### Version of Awkward Array

2.2.4

### Description and code to reproduce

While trying out awkward + jax I ran into the following:
```python
import jax
import awkward as ak

ak.jax.register_and_check()

a = ak.Array([[1.0, 2, 3], [5, 6]], backend="jax")

def f(x):
    return ak.sum(ak.sum(x) * x)

f(a)
jax.grad(f)(a)
```
results in
```python
[...]
TypeError: Jax tracers cannot be used with `ak.from_arraylib`

This error occurred while calling

    numpy.multiply.__call__(
        JVPTracer-instance
        <Array [[...], [...]] type='2 * var * float32'>
    )
```
This seems to originate from the `ak.sum(x) * x` piece.

A corresponding `numpy` version (without any `awkward`) works fine:
```python
import jax
import numpy as np

a = np.array([1, 2, 3], dtype=float)

def f(x):
    return np.sum(np.sum(x) * x)

f(a)
jax.grad(f)(a)
```

Relevant library versions:
```
awkward                       2.2.4
awkward-cpp                   17
jax                           0.4.13
jaxlib                        0.4.13
numpy                         1.25.0
```



Discussion:
@Saransh-cpp, when you have access to the Awkward repo (see https://github.com/scikit-hep/awkward/invitations), please assign yourself to this issue. Possible outcomes for this issue include:

* testing it and realizing that it's not broken anymore, and closing it without a fix
* having some good justification for why it's not important anymore, and closing it without a fix
* fixing it

I'm just organizing responsibilities at this point.
I tried removing the intentional error and everything seems to work -
```python
In [1]: import jax
   ...: import numpy as np
   ...: 
   ...: a = np.array([1, 2, 3], dtype=float)
   ...: 
   ...: def f(x):
   ...:     return np.sum(np.sum(x) * x)
   ...: 
   ...: f(a)
   ...: jax.grad(f)(a)
Out[1]: Array([12., 12., 12.], dtype=float32)

In [2]: import jax
   ...: import awkward as ak
   ...: 
   ...: ak.jax.register_and_check()
   ...: 
   ...: a = ak.Array([[1.0, 2, 3], [5, 6]], backend="jax")
   ...: 
   ...: def f(x):
   ...:     return ak.sum(ak.sum(x) * x)
   ...: 
   ...: f(a)
   ...: jax.grad(f)(a)
Out[2]: <Array [[34.0, 34.0, 34.0], [34.0, 34.0]] type='2 * var * float32'>
``` 

Surprisingly, all the tests pass too. @agoose77, given that you tried this above but ended up closing the PR, could you please let me know why this intentional error is required? Thanks!
@Saransh-cpp at the time, I recall having thoughts that seeing tracers at that point was a symptom of a problem that we needed to fix rather than just removing the error. However, my memory of this is not brilliant.

I would caution that our tests aren't super detailed here, but nevertheless support removing the guard until you find something that tells you otherwise! Maybe I mis-remembered, or it's actually not a problem after all ...
Oh, I see, thanks! I'll create a PR removing the error to see if everything goes right.

Resolving patch:
diff --git a/src/awkward/_layout.py b/src/awkward/_layout.py
--- a/src/awkward/_layout.py
+++ b/src/awkward/_layout.py
@@ -13,7 +13,6 @@
 from awkward._backends.numpy import NumpyBackend
 from awkward._behavior import behavior_of
 from awkward._nplikes.dispatch import nplike_of_obj
-from awkward._nplikes.jax import Jax
 from awkward._nplikes.numpy import Numpy
 from awkward._nplikes.numpy_like import NumpyMetadata
 from awkward._typing import TYPE_CHECKING, Any, NamedTuple, Self, TypeVar
@@ -266,10 +265,6 @@ def from_arraylib(array, regulararray, recordarray):
     nplike = nplike_of_obj(array)
 
     def recurse(array, mask=None):
-        cls = type(array)
-        if Jax.is_tracer_type(cls):
-            raise TypeError("Jax tracers cannot be used with `ak.from_arraylib`")
-
         if regulararray and len(array.shape) > 1:
             new_shape = (-1,) + array.shape[2:]
             return RegularArray(


Test patch:
diff --git a/tests/test_2556_jax_tracer_error.py b/tests/test_2556_jax_tracer_error.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2556_jax_tracer_error.py
@@ -0,0 +1,20 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+from __future__ import annotations
+
+import pytest
+
+import awkward as ak
+
+jax = pytest.importorskip("jax")
+
+
+def test():
+    ak.jax.register_and_check()
+
+    arr = ak.Array([[1.0, 2, 3], [5, 6]], backend="jax")
+    grad_arr = ak.Array([[34.0, 34.0, 34.0], [34.0, 34.0]], backend="jax")
+
+    def f(x):
+        return ak.sum(ak.sum(x) * x)
+
+    assert ak.all(grad_arr == jax.grad(f)(arr))
diff --git a/tests/test_2637_jax_tracer_error.py b/tests/test_2637_jax_tracer_error.py
new file mode 100644
--- /dev/null
+++ b/tests/test_2637_jax_tracer_error.py
@@ -0,0 +1,37 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
+from __future__ import annotations
+
+import pytest
+
+import awkward as ak
+
+jax = pytest.importorskip("jax")
+
+
+def test():
+    ak.jax.register_and_check()
+
+    jets = ak.Array(
+        [
+            [
+                {"pt": 1.0, "eta": 1.1, "phi": 0.1, "mass": 0.01},
+                {"pt": 2, "eta": 2.2, "phi": 0.2, "mass": 0.02},
+            ],
+            [
+                {"pt": 4.0, "eta": 4.4, "phi": 0.4, "mass": 0.04},
+                {"pt": 5.0, "eta": 5.5, "phi": 0.5, "mass": 0.05},
+                {"pt": 6.0, "eta": 6.6, "phi": 0.6, "mass": 0.06},
+            ],
+        ],
+        backend="jax",
+    )
+
+    def correct_jets(jets, alpha):
+        new_pt = jets["pt"] + 25.0 * alpha
+        jets["pt"] = new_pt
+        return ak.sum(jets["pt"])
+
+    val, grad = jax.value_and_grad(correct_jets, argnums=1)(jets, 0.1)
+
+    assert val == 30.5
+    assert grad == 125.0