Merge pull request #79 from e3nn/get_parameter

Fix documentation build errors.

Author: ameya98

e3nn_jax/_src/basic.py

@@ -648,8 +648,11 @@ def normal(
             raise ValueError("Normalization needs to be 'norm' or 'component'")
 
 
-def where(mask: jax.Array, x: e3nn.IrrepsArray, y: e3nn.IrrepsArray):
-    """Selects elements from `x` or `y`, depending on `mask`.
+def where(
+    mask: jax.Array, x: e3nn.IrrepsArray, y: e3nn.IrrepsArray
+) -> e3nn.IrrepsArray:
+    """
+    Selects elements from `x` or `y`, depending on `mask`.
 
     Equivalent to:
         >>> e3nn.IrrepsArray(x.irreps, jnp.where(mask, x.array, y.array))

e3nn_jax/_src/batchnorm/bn_flax.py

@@ -25,7 +25,7 @@ class BatchNorm(nn.Module):
 
     Args:
         use_running_average: if True, the statistics stored in batch_stats will be
-      used instead of computing the batch statistics on the input.
+            used instead of computing the batch statistics on the input.
         eps (float): epsilon for numerical stability, has to be between 0 and 1.
             the field norm is transformed to ``(1 - eps) * norm + eps``
             leading to a slower convergence toward norm 1.

e3nn_jax/_src/irreps_array.py

@@ -38,7 +38,7 @@ def _is_none_slice(x):
     return isinstance(x, slice) and x == slice(None)
 
 
-@attrs(frozen=True, init=True, repr=False)
+@attrs(frozen=True, init=True, repr=False, cmp=False)
 class IrrepsArray:
     r"""Array with a representation of rotations.