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,