u1d-16bit-math (#188)

Author: mariogeiger

cuequivariance_jax/cuequivariance_jax/segmented_polynomials/segmented_polynomial.py

@@ -85,7 +85,7 @@ def segmented_polynomial(
 
             - ``"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.
+            - ``"uniform_1d"``: String dtype names (e.g., ``"float32"``, ``"float64"``, ``"float16"``, ``"bfloat16"``).
             - ``"indexed_linear"``: CUBLAS compute type strings such as ``"CUBLAS_COMPUTE_32F"``, ``"CUBLAS_COMPUTE_32F_FAST_TF32"``,
               ``"CUBLAS_COMPUTE_32F_PEDANTIC"``, ``"CUBLAS_COMPUTE_64F"``, etc.
 

cuequivariance_jax/cuequivariance_jax/segmented_polynomials/segmented_polynomial_uniform_1d.py

@@ -152,7 +152,7 @@ def fn(op, d: cue.SegmentedTensorProduct):
         raise ValueError(f"Buffer shapes not compatible {[b.shape for b in buffers]}")
 
     if math_dtype is not None:
-        supported_dtypes = {"float32", "float64"}
+        supported_dtypes = {"float32", "float64", "float16", "bfloat16"}
         if math_dtype not in supported_dtypes:
             raise ValueError(
                 f"method='uniform_1d' only supports math_dtype equal to {supported_dtypes}, got '{math_dtype}'."