More comments

README.md

@@ -59,7 +59,7 @@ Where each group supports:
   <code>jaxlie.<strong>manifold.\*</strong></code>).
 - Compatibility with standard JAX function transformations. (see
   [./examples/vmap_example.py](./examples/vmap_example.py))
-- Broadcasting.
+- Broadcasting for leading axes.
 - (Un)flattening as pytree nodes.
 - Serialization using [flax](https://github.com/google/flax).
 

jaxlie/_se2.py

@@ -18,7 +18,8 @@
 )
 @jdc.pytree_dataclass
 class SE2(_base.SEBase[SO2]):
-    """Special Euclidean group for proper rigid transforms in 2D.
+    """Special Euclidean group for proper rigid transforms in 2D. Broadcasting
+    rules are the same as for numpy.
 
     Internal parameterization is `(cos, sin, x, y)`. Tangent parameterization is `(vx,
     vy, omega)`.

jaxlie/_se3.py

@@ -31,7 +31,8 @@ def _skew(omega: hints.Array) -> jax.Array:
 )
 @jdc.pytree_dataclass
 class SE3(_base.SEBase[SO3]):
-    """Special Euclidean group for proper rigid transforms in 3D.
+    """Special Euclidean group for proper rigid transforms in 3D. Broadcasting
+    rules are the same as for numpy.
 
     Internal parameterization is `(qw, qx, qy, qz, x, y, z)`. Tangent parameterization
     is `(vx, vy, vz, omega_x, omega_y, omega_z)`.

jaxlie/_so2.py

@@ -19,7 +19,8 @@
 )
 @jdc.pytree_dataclass
 class SO2(_base.SOBase):
-    """Special orthogonal group for 2D rotations.
+    """Special orthogonal group for 2D rotations. Broadcasting rules are the
+    same as for `numpy`.
 
     Internal parameterization is `(cos, sin)`. Tangent parameterization is `(omega,)`.
     """

jaxlie/_so3.py

@@ -19,7 +19,8 @@
 )
 @jdc.pytree_dataclass
 class SO3(_base.SOBase):
-    """Special orthogonal group for 3D rotations.
+    """Special orthogonal group for 3D rotations. Broadcasting rules are the same as
+    for numpy.
 
     Internal parameterization is `(qw, qx, qy, qz)`. Tangent parameterization is
     `(omega_x, omega_y, omega_z)`.