add sharding notes (#197)

Author: mariogeiger

cuequivariance_jax/cuequivariance_jax/segmented_polynomials/segmented_polynomial.py

@@ -146,6 +146,14 @@ def segmented_polynomial(
         ... )
         >>> y.shape
         (10, 80)
+
+
+    .. note::
+        This operation uses a custom CUDA kernel for performance. When using this function
+        on multiple devices, manual sharding is required to achieve proper performance.
+        Without explicit sharding, performance will be significantly degraded. See
+        `JAX shard_map documentation <https://docs.jax.dev/en/latest/notebooks/shard_map.html>`_
+        for details on manual parallelism.
     """
 
     if method == "":

cuequivariance_jax/cuequivariance_jax/triangle/_triangle_attention.py

@@ -61,6 +61,13 @@ def triangle_attention(
 
     where :math:`Q`, :math:`K`, and :math:`V` are the query, key, and value tensors,
     :math:`M` is the mask bias, and :math:`T` is the triangle bias.
+
+    .. note::
+        This operation uses a custom CUDA kernel for performance. When using this function
+        on multiple devices, manual sharding is required to achieve proper performance.
+        Without explicit sharding, performance will be significantly degraded. See
+        `JAX shard_map documentation <https://docs.jax.dev/en/latest/notebooks/shard_map.html>`_
+        for details on manual parallelism.
     """
     return triangle_attention_custom_vjp(
         q, k, v, bias, mask, scale=scale, precision=precision

cuequivariance_jax/cuequivariance_jax/triangle/_triangle_multiplicative_update.py

@@ -239,6 +239,13 @@ def triangle_multiplicative_update(
         ... )
         >>> print(output_256.shape)
         (2, 3, 128, 128, 256)
+
+    .. note::
+        This operation uses a custom CUDA kernel for performance. When using this function
+        on multiple devices, manual sharding is required to achieve proper performance.
+        Without explicit sharding, performance will be significantly degraded. See
+        `JAX shard_map documentation <https://docs.jax.dev/en/latest/notebooks/shard_map.html>`_
+        for details on manual parallelism.
     """
     # Input validation
     if direction not in ["outgoing", "incoming"]: