Fix.

e3nn · Aug 14, 2024 · 129706a · 129706a
1 parent 0c7fe1c
commit 129706a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/e3nn_jax/_src/basic.py b/e3nn_jax/_src/basic.py
@@ -648,7 +648,9 @@ def normal(
             raise ValueError("Normalization needs to be 'norm' or 'component'")
 
 
-def where(mask: jax.Array, x: e3nn.IrrepsArray, y: e3nn.IrrepsArray) -> e3nn.IrrepsArray:
+def where(
+    mask: jax.Array, x: e3nn.IrrepsArray, y: e3nn.IrrepsArray
+) -> e3nn.IrrepsArray:
     """
     Selects elements from `x` or `y`, depending on `mask`.