perf: reduce python overhead for awkward backend (#554)

* awkward backend: reduce python overhead by applying multiple operations in one single broadcasting traversal

* add a fast path for noops, otherwise ak.transform does a non-negligible overhead...

Author: pfackeldey

src/vector/__init__.py

@@ -76,7 +76,7 @@ def _import_awkward() -> None:
     if not typing.TYPE_CHECKING:
         VectorAwkward = None
 else:
-    from vector.backends.awkward import VectorAwkward
+    from vector.backends.awkward import VectorAwkward, awkward_transform
 
 try:
     import sympy  # type: ignore[import-untyped]
@@ -143,6 +143,7 @@ def _import_awkward() -> None:
     "arr",
     "array",
     "awk",
+    "awkward_transform",
     "dim",
     "obj",
     "register_awkward",

src/vector/_compute/lorentz/Et.py

@@ -114,7 +114,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/Et2.py

@@ -116,7 +116,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/Mt.py

@@ -110,7 +110,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/Mt2.py

@@ -111,7 +111,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/add.py

@@ -201,9 +201,10 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/beta.py

@@ -153,7 +153,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/boostX_beta.py

@@ -243,7 +243,7 @@ def dispatch(beta: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 beta,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/boostX_gamma.py

@@ -243,7 +243,7 @@ def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 gamma,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/boostY_beta.py

@@ -243,7 +243,7 @@ def dispatch(beta: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 beta,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/boostY_gamma.py

@@ -243,7 +243,7 @@ def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 gamma,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/boostZ_beta.py

@@ -218,7 +218,7 @@ def dispatch(beta: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 beta,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/boostZ_gamma.py

@@ -218,7 +218,7 @@ def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 gamma,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/boost_beta3.py

@@ -391,9 +391,10 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/boost_p4.py

@@ -781,9 +781,10 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/deltaRapidityPhi.py

@@ -110,9 +110,10 @@ def dispatch(
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/deltaRapidityPhi2.py

@@ -106,9 +106,10 @@ def dispatch(
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/dot.py

@@ -155,9 +155,10 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/equal.py

@@ -168,9 +168,10 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/gamma.py

@@ -153,7 +153,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/is_lightlike.py

@@ -68,7 +68,7 @@ def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 tolerance,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/is_spacelike.py

@@ -66,7 +66,7 @@ def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 tolerance,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/is_timelike.py

@@ -66,7 +66,7 @@ def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 tolerance,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/isclose.py

@@ -221,9 +221,10 @@ def dispatch(
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 rtol,
                 atol,

src/vector/_compute/lorentz/not_equal.py

@@ -168,9 +168,10 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/rapidity.py

@@ -127,7 +127,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/scale.py

@@ -181,7 +181,7 @@ def dispatch(factor: typing.Any, v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 factor,
                 *v.azimuthal.elements,

src/vector/_compute/lorentz/subtract.py

@@ -201,9 +201,10 @@ def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any:
         ),
     )
     with numpy.errstate(all="ignore"):
-        return _handler_of(v1, v2)._wrap_result(
+        handler = _handler_of(v1, v2)
+        return handler._wrap_result(
             _flavor_of(v1, v2),
-            function(
+            handler._wrap_dispatched_function(function)(
                 _lib_of(v1, v2),
                 *v1.azimuthal.elements,
                 *v1.longitudinal.elements,

src/vector/_compute/lorentz/t.py

@@ -37,6 +37,9 @@ def xy_z_t(lib, x, y, z, t):
     return t
 
 
+xy_z_t.__awkward_transform_allowed__ = False  # type:ignore[attr-defined]
+
+
 def xy_z_tau(lib, x, y, z, tau):
     return lib.sqrt(t2.xy_z_tau(lib, x, y, z, tau))
 
@@ -45,6 +48,9 @@ def xy_theta_t(lib, x, y, theta, t):
     return t
 
 
+xy_theta_t.__awkward_transform_allowed__ = False  # type:ignore[attr-defined]
+
+
 def xy_theta_tau(lib, x, y, theta, tau):
     return lib.sqrt(t2.xy_theta_tau(lib, x, y, theta, tau))
 
@@ -53,6 +59,9 @@ def xy_eta_t(lib, x, y, eta, t):
     return t
 
 
+xy_eta_t.__awkward_transform_allowed__ = False  # type:ignore[attr-defined]
+
+
 def xy_eta_tau(lib, x, y, eta, tau):
     return lib.sqrt(t2.xy_eta_tau(lib, x, y, eta, tau))
 
@@ -61,6 +70,9 @@ def rhophi_z_t(lib, rho, phi, z, t):
     return t
 
 
+rhophi_z_t.__awkward_transform_allowed__ = False  # type:ignore[attr-defined]
+
+
 def rhophi_z_tau(lib, rho, phi, z, tau):
     return lib.sqrt(t2.rhophi_z_tau(lib, rho, phi, z, tau))
 
@@ -69,6 +81,9 @@ def rhophi_theta_t(lib, rho, phi, theta, t):
     return t
 
 
+rhophi_theta_t.__awkward_transform_allowed__ = False  # type:ignore[attr-defined]
+
+
 def rhophi_theta_tau(lib, rho, phi, theta, tau):
     return lib.sqrt(t2.rhophi_theta_tau(lib, rho, phi, theta, tau))
 
@@ -77,6 +92,9 @@ def rhophi_eta_t(lib, rho, phi, eta, t):
     return t
 
 
+rhophi_eta_t.__awkward_transform_allowed__ = False  # type:ignore[attr-defined]
+
+
 def rhophi_eta_tau(lib, rho, phi, eta, tau):
     return lib.sqrt(t2.rhophi_eta_tau(lib, rho, phi, eta, tau))
 
@@ -110,7 +128,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,

src/vector/_compute/lorentz/t2.py

@@ -125,7 +125,7 @@ def dispatch(v: typing.Any) -> typing.Any:
     with numpy.errstate(all="ignore"):
         return v._wrap_result(
             _flavor_of(v),
-            function(
+            v._wrap_dispatched_function(function)(
                 v.lib,
                 *v.azimuthal.elements,
                 *v.longitudinal.elements,