Array API version of norm_angle

dfki-ric · Aug 23, 2023 · aa862c9 · aa862c9
1 parent dd8b85c
commit aa862c9
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/pytransform3d/rotations/_utils.py b/pytransform3d/rotations/_utils.py
@@ -72,8 +72,9 @@ def norm_angle(a):
     a_norm : float or array-like, shape (n,)
         Normalized angle(s) in radians
     """
-    a = np.asarray(a, dtype=np.float64)
-    return (a - (np.ceil((a + np.pi) / two_pi) - 1.0) * two_pi)
+    xp = array_namespace(a)
+    a = xp.asarray(a, dtype=xp.float64)
+    return a - (xp.ceil((a + xp.pi) / two_pi) - 1.0) * two_pi
 
 
 def norm_axis_angle(a):