fix: correctly infer the numpy-like module

pyproject.toml

@@ -59,6 +59,7 @@ optional-dependencies.dev = [
   "pytest-cov>=3",
   "pytest-doctestplus",
   "sympy",
+  "jax",
 ]
 optional-dependencies.docs = [
   "awkward>=2",
@@ -85,6 +86,7 @@ optional-dependencies.test = [
   "pytest-doctestplus",
 ]
 optional-dependencies.test-extras = [
+  "jax",
   "dask_awkward",
   "spark-parser",
   'uncompyle6; python_version == "3.8"',

src/vector/backends/awkward.py

@@ -649,11 +649,10 @@ class VectorAwkward:
 
     @property
     def lib(self):  # type:ignore[no-untyped-def]
-        if (
-            nplike := self.layout.backend.nplike  # type:ignore[attr-defined]
-        ) is ak._nplikes.typetracer.TypeTracer.instance():
+        nplike = self.layout.backend.nplike  # type:ignore[attr-defined]
+        if nplike is ak._nplikes.typetracer.TypeTracer.instance():
             return _lib(module=numpy, nplike=nplike)
-        return numpy
+        return _lib(module=nplike._module, nplike=nplike)
 
     def _wrap_result(
         self: AwkwardProtocol,

tests/test_issues.py

@@ -170,3 +170,18 @@ def test_issue_463():
     for transform in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta":
         trv = getattr(v, "to_" + transform)()
         assert trv.deltaangle(trv) == 0.0
+
+
+def test_issue_621():
+    _ = pytest.importorskip("jax")
+    ak = pytest.importorskip("awkward")
+    vector.register_awkward()
+    ak.jax.register_and_check()
+
+    a = b = ak.to_backend(
+        ak.zip({"x": [1], "y": [1], "z": [1], "t": [1]}, with_name="Momentum4D"), "jax"
+    )
+
+    # some computation that involves broadcast_and_apply in awkward
+    # enough to check if it computes at all
+    assert (a + b).mass