add missing asarray in axis_angle_to_matrix

e3nn · Oct 27, 2023 · 2375789 · 2375789
1 parent 6f3a826
commit 2375789
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/e3nn_jax/_src/rotation.py b/e3nn_jax/_src/rotation.py
@@ -590,6 +590,7 @@ def axis_angle_to_matrix(axis, angle):
     Returns:
         `jax.numpy.ndarray`: array of shape :math:`(..., 3, 3)`
     """
+    angle = jnp.asarray(angle)
     axis, angle = jnp.broadcast_arrays(axis, angle[..., None])
     alpha, beta = xyz_to_angles(axis)
     R = angles_to_matrix(alpha, beta, jnp.zeros_like(beta))