Revert "[JAX] expand math_dtype argument (#153)"

This reverts commit 9b6788b.

Author: mariogeiger

cuequivariance_jax/cuequivariance_jax/equivariant_polynomial.py

@@ -29,9 +29,9 @@ def equivariant_polynomial(
     indices: None | list[None | jax.Array | tuple[jax.Array | slice]] = None,
     *,
     method: str = "",
-    math_dtype: str | None = None,
+    math_dtype: jnp.dtype | None = None,
     name: str | None = None,
-    precision: jax.lax.Precision = "undefined",
+    precision: jax.lax.Precision = jax.lax.Precision.HIGHEST,
 ) -> list[cuex.RepArray] | cuex.RepArray:
     """Compute an equivariant polynomial.
 
@@ -50,8 +50,10 @@ def equivariant_polynomial(
             total number of operands (inputs + outputs). Use None for unindexed
             operands. Defaults to None. Note that indices are not supported for all methods.
         method: Method to use for computation. See :func:`cuex.segmented_polynomial <cuequivariance_jax.segmented_polynomial>` for available methods.
-        math_dtype: See :func:`cuex.segmented_polynomial <cuequivariance_jax.segmented_polynomial>` for supported options.
+        math_dtype: Data type for computational operations. If None, automatically
+            determined from input types. Defaults to None.
         name: Optional name for the operation. Defaults to None.
+        precision: The precision to use for the computation. Defaults to HIGHEST. Note that precision is not supported for all methods.
 
     Returns:
         :class:`cuex.RepArray <cuequivariance_jax.RepArray>` or list of :class:`cuex.RepArray <cuequivariance_jax.RepArray>`
@@ -102,11 +104,6 @@ def equivariant_polynomial(
     if name is None:
         name = "equivariant_polynomial"
 
-    if precision != "undefined":
-        raise ValueError(
-            "precision is not anymore supported. Please use math_dtype instead."
-        )
-
     if len(inputs) != poly.num_inputs:
         raise ValueError(
             f"Unexpected number of inputs. Expected {poly.num_inputs}, got {len(inputs)}."
@@ -188,6 +185,7 @@ def equivariant_polynomial(
         math_dtype=math_dtype,
         name=name,
         method=method,
+        precision=precision,
     )
     outputs = [cuex.RepArray(rep, x) for rep, x in zip(poly.outputs, outputs)]
 

cuequivariance_jax/cuequivariance_jax/experimental/indexed_linear.py

@@ -14,18 +14,84 @@
 # limitations under the License.
 import jax
 import jax.lax
+import jax.numpy as jnp
 
 import cuequivariance as cue
+import cuequivariance_jax as cuex
 
 
 def indexed_linear(
     poly: cue.SegmentedPolynomial,
     counts: jax.Array,
     w: jax.Array,
     x: jax.Array,
-    math_dtype: str | None = None,
+    math_dtype: jnp.dtype | None = None,
     method: str = "indexed_linear",
 ) -> jax.Array:
-    raise NotImplementedError(
-        "Use cuex.segmented_polynomial(..., method='indexed_linear') instead."
+    """Linear layer with different weights for different parts of the input.
+
+    Args:
+        poly: The polynomial descriptor. Only works for descriptors of a linear layer.
+        counts: Number of elements in each partition. Shape (C,).
+        w: Weights of the linear layer. Shape (C, num_weights).
+        x: Input data. Shape (Z, num_inputs). Z is equal to the sum of counts.
+        math_dtype: Data type for computational operations. If
+            None, automatically determined from input types. Defaults to None.
+    Returns:
+        Output data. Shape (Z, num_outputs).
+
+    Examples:
+        This example demonstrates using indexed_linear for a batch of inputs with
+        different species:
+
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> import cuequivariance as cue
+        >>> import cuequivariance_jax as cuex
+        >>>
+        >>> # Define problem parameters
+        >>> num_species_total = 3  # Total number of different species
+        >>> batch_size = 10        # Number of samples in batch
+        >>> input_dim = 8          # Input feature dimension
+        >>> output_dim = 16        # Output feature dimension
+        >>> dtype = jnp.float32
+        >>>
+        >>> # Define how many elements belong to each species
+        >>> num_species = jnp.array([3, 4, 3], dtype=jnp.int32)  # Sum equals batch_size
+        >>>
+        >>> # Generate random input data
+        >>> input_array = jax.random.normal(jax.random.key(0), (batch_size, input_dim), dtype)
+        >>>
+        >>> # Define irreps for input and output features
+        >>> input_irreps = cue.Irreps(cue.O3, f"{input_dim}x0e")   # Scalar features
+        >>> output_irreps = cue.Irreps(cue.O3, f"{output_dim}x0e") # Scalar features
+        >>>
+        >>> # Create a linear descriptor
+        >>> e = cue.descriptors.linear(input_irreps, output_irreps)
+        >>>
+        >>> # Generate weights for each species
+        >>> w = jax.random.normal(jax.random.key(1), (num_species_total, e.inputs[0].dim), dtype)
+        >>>
+        >>> # Apply the indexed linear layer
+        >>> result = cuex.experimental.indexed_linear(e.polynomial, num_species, w, input_array)
+        >>>
+        >>> # Verify output shape
+        >>> assert result.shape == (batch_size, output_dim)
+    """
+    assert poly.num_inputs == 2
+    assert poly.num_outputs == 1
+
+    (C, _) = w.shape
+    (Z, _) = x.shape
+    assert counts.shape == (C,)
+
+    y = jax.ShapeDtypeStruct((Z, poly.outputs[0].size), x.dtype)
+    [y] = cuex.segmented_polynomial(
+        poly,
+        [w, x],
+        [y],
+        [cuex.Repeats(counts), None, None],
+        math_dtype=math_dtype,
+        method=method,
     )
+    return y

cuequivariance_jax/cuequivariance_jax/segmented_polynomials/segmented_polynomial.py

@@ -54,9 +54,9 @@ def segmented_polynomial(
     indices: None | list[None | jax.Array | tuple[jax.Array | slice]] = None,
     *,
     method: str = "",
-    math_dtype: str | None = None,
+    math_dtype: jnp.dtype | None = None,
     name: str | None = None,
-    precision: jax.lax.Precision = "undefined",
+    precision: jax.lax.Precision = jax.lax.Precision.HIGHEST,
 ) -> list[jax.Array]:
     """Compute a segmented polynomial.
 
@@ -79,17 +79,13 @@ def segmented_polynomial(
 
             .. note::
                The ``"fused_tp"`` method is only available in the PyTorch implementation.
-        math_dtype: Data type for computational operations. If None, automatically determined from input types. Defaults to None.
-
-            Supported options vary by method:
-
-            - ``"naive"``: String dtype names (e.g., ``"float32"``, ``"float64"``, ``"float16"``, ``"bfloat16"``).
-              Also supports ``"tensor_float32"`` for TensorFloat-32 mode.
-            - ``"uniform_1d"``: String values ``"float32"`` or ``"float64"`` only.
-            - ``"indexed_linear"``: CUBLAS compute type strings such as ``"CUBLAS_COMPUTE_32F"``, ``"CUBLAS_COMPUTE_32F_FAST_TF32"``,
-              ``"CUBLAS_COMPUTE_32F_PEDANTIC"``, ``"CUBLAS_COMPUTE_64F"``, etc.
+        math_dtype: Data type for computational operations. If None, automatically
+            determined from input types, defaulting to float32 if no float64 inputs
+            are present.
 
         name: Optional name for the operation.
+        precision: The precision to use for the computation. Defaults to HIGHEST.
+            Note that precision is only supported for the ``"naive"`` method.
 
     Returns:
         List of JAX arrays containing the computed polynomial outputs.
@@ -163,15 +159,6 @@ def segmented_polynomial(
     if name is None:
         name = "segmented_polynomial"
 
-    if math_dtype is not None and not isinstance(math_dtype, str):
-        math_dtype = jnp.dtype(math_dtype).name
-    assert isinstance(math_dtype, str) or math_dtype is None
-
-    if precision != "undefined":
-        raise ValueError(
-            "precision is not anymore supported. Please use math_dtype instead."
-        )
-
     assert len(inputs) == polynomial.num_inputs
     assert len(outputs_shape_dtype) == polynomial.num_outputs
 
@@ -281,6 +268,10 @@ def fn(x, n: int):
             index_configuration.append(bi)
             index_mode.append(im)
 
+    # Set default math_dtype
+    if math_dtype is None:
+        math_dtype = jnp.result_type(*io_buffers)
+
     # Execute the polynomial
     kwargs = dict(
         inputs=io_buffers[: polynomial.num_inputs],
@@ -291,6 +282,7 @@ def fn(x, n: int):
         polynomial=polynomial,
         math_dtype=math_dtype,
         name=name,
+        precision=precision,
     )
 
     outputs = segmented_polynomial_prim(**kwargs, method=method)
@@ -352,9 +344,10 @@ def segmented_polynomial_prim(
     index_configuration: list[list[int]],  # maps: buffer index -> unique indices index
     index_mode: list[list[IndexingMode]],  # shared, batched, indexed, repeated
     polynomial: cue.SegmentedPolynomial,
-    math_dtype: str | None,
+    math_dtype: jnp.dtype,
     name: str,
     method: str,
+    precision: jax.lax.Precision,
     return_none_if_empty: bool = False,
 ) -> tuple[jax.Array, ...]:  # output buffers
     """
@@ -383,9 +376,10 @@ def segmented_polynomial_prim(
             x for x, used in zip(outputs_shape_dtype, used_outputs) if used
         ),
         polynomial=polynomial.filter_keep_operands(used_inputs + used_outputs),
-        math_dtype=math_dtype,
+        math_dtype=jnp.dtype(math_dtype),
         name=str(name),
         method=method,
+        precision=precision,
     )
 
     if return_none_if_empty:
@@ -441,9 +435,10 @@ def segmented_polynomial_abstract_eval(
     index_mode: tuple[tuple[IndexingMode, ...], ...],
     outputs_shape_dtype: tuple[jax.ShapeDtypeStruct, ...],
     polynomial: cue.SegmentedPolynomial,
-    math_dtype: str | None,
+    math_dtype: jnp.dtype,
     name: str,
     method: str,
+    precision: jax.lax.Precision,
 ) -> tuple[jax.core.ShapedArray, ...]:
     return tuple(
         jax.core.ShapedArray(out.shape, out.dtype) for out in outputs_shape_dtype
@@ -457,9 +452,10 @@ def segmented_polynomial_impl(
     index_mode: tuple[tuple[IndexingMode, ...], ...],
     outputs_shape_dtype: tuple[jax.ShapeDtypeStruct, ...],
     polynomial: cue.SegmentedPolynomial,
-    math_dtype: str | None,
+    math_dtype: jnp.dtype,
     name: str,
     method: str,
+    precision: jax.lax.Precision,
 ) -> tuple[jax.Array, ...]:
     num_inputs = len(index_configuration) - len(outputs_shape_dtype)
     inputs, indices = inputs_and_indices[:num_inputs], inputs_and_indices[num_inputs:]
@@ -510,14 +506,21 @@ def segmented_polynomial_impl(
             raise ValueError(
                 "IndexingMode.REPEATED is only supported with 'naive' or 'indexed_linear' methods."
             )
+    if precision != jax.lax.Precision.HIGHEST:
+        if method not in ("naive",):
+            raise ValueError(
+                f"Precision {precision} is only supported with 'naive' method."
+            )
 
     match method:
         case "naive":
-            return execute_naive(**kwargs, index_mode=index_mode)
+            return execute_naive(**kwargs, index_mode=index_mode, precision=precision)
         case "uniform_1d":
             return execute_uniform_1d(**kwargs)
         case "indexed_linear":
-            return execute_indexed_linear(**kwargs, index_mode=index_mode)
+            return execute_indexed_linear(
+                **kwargs, index_mode=index_mode, precision=precision
+            )
 
 
 def segmented_polynomial_jvp(
@@ -528,9 +531,10 @@ def segmented_polynomial_jvp(
     index_mode: tuple[tuple[IndexingMode, ...], ...],
     outputs_shape_dtype: tuple[jax.ShapeDtypeStruct, ...],
     polynomial: cue.SegmentedPolynomial,
-    math_dtype: str | None,
+    math_dtype: jnp.dtype,
     name: str,
     method: str,
+    precision: jax.lax.Precision,
 ) -> tuple[tuple[jax.Array, ...], tuple[jax.Array | ad.Zero, ...]]:
     num_inputs = len(index_configuration) - len(outputs_shape_dtype)
 
@@ -552,6 +556,7 @@ def segmented_polynomial_jvp(
         math_dtype,
         name,
         method=method,
+        precision=precision,
     )
 
     jvp_poly, _ = polynomial.jvp([not isinstance(t, ad.Zero) for t in tangents])
@@ -576,6 +581,7 @@ def segmented_polynomial_jvp(
         + "_jvp"
         + "".join("0" if isinstance(t, ad.Zero) else "1" for t in tangents),
         method=method,
+        precision=precision,
     )
 
     return out_primals, out_tangents
@@ -588,9 +594,10 @@ def segmented_polynomial_transpose(
     index_mode: tuple[tuple[IndexingMode, ...], ...],
     outputs_shape_dtype: tuple[jax.ShapeDtypeStruct, ...],
     polynomial: cue.SegmentedPolynomial,
-    math_dtype: str | None,
+    math_dtype: jnp.dtype,
     name: str,
     method: str,
+    precision: jax.lax.Precision,
 ) -> tuple[jax.Array | ad.Zero | None, ...]:
     num_inputs = len(index_configuration) - len(outputs_shape_dtype)
     inputs, indices = inputs_and_indices[:num_inputs], inputs_and_indices[num_inputs:]
@@ -632,6 +639,7 @@ def segmented_polynomial_transpose(
         math_dtype,
         name + "_T",
         method=method,
+        precision=precision,
         return_none_if_empty=True,
     )
 
@@ -652,9 +660,10 @@ def segmented_polynomial_batching(
     index_mode: tuple[tuple[IndexingMode, ...], ...],
     outputs_shape_dtype: tuple[jax.ShapeDtypeStruct, ...],
     polynomial: cue.SegmentedPolynomial,
-    math_dtype: str | None,
+    math_dtype: jnp.dtype,
     name: str,
     method: str,
+    precision: jax.lax.Precision,
 ) -> tuple[tuple[jax.Array, ...], tuple[int, ...]]:
     # Add a new batch axis in the first dimension
     def prepare(input: jax.Array, axis: int | None) -> jax.Array:
@@ -695,6 +704,7 @@ def prepare(input: jax.Array, axis: int | None) -> jax.Array:
         math_dtype=math_dtype,
         name=name + "_batching",
         method=method,
+        precision=precision,
     )
     return outputs, (0,) * len(outputs)