feat: dtype mixed-precision override for einsum (#22)

Author: scttfrdmn

CHANGELOG.md

@@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- **`dtype` mixed-precision override for `einsum`** (#22) — accepts string
+  aliases (`"bf16"`, `"fp16"`, `"f32"`) or `torch.dtype` instances. When set,
+  all operands are cast to the requested dtype before contracting and the
+  result is returned in that dtype. Matches Neuron SDK autocast recommendations;
+  use `dtype="bf16"` to route fp32 models through the NKI bf16 matmul path
+  without changing the model's weight dtype.
+
 ## [0.3.0] — 2026-04-16
 
 ### Added

tests/test_einsum.py

@@ -246,3 +246,48 @@ def test_defaults_unchanged(self):
         result = trntensor.einsum("ij,jk->ik", A, B, alpha=1.0, beta=0.0, out=None)
         expected = A @ B
         np.testing.assert_allclose(result.numpy(), expected.numpy(), atol=1e-5)
+
+
+class TestDtype:
+    """Tests for the dtype= mixed-precision override (#22)."""
+
+    def test_dtype_bf16_result(self):
+        A = torch.randn(4, 3)
+        B = torch.randn(3, 5)
+        result = trntensor.einsum("ij,jk->ik", A, B, dtype="bf16")
+        assert result.dtype == torch.bfloat16
+
+    def test_dtype_fp16_result(self):
+        A = torch.randn(4, 3)
+        B = torch.randn(3, 5)
+        result = trntensor.einsum("ij,jk->ik", A, B, dtype="fp16")
+        assert result.dtype == torch.float16
+
+    def test_dtype_torch_type(self):
+        """Passing a torch.dtype directly is equivalent to the string alias."""
+        A = torch.randn(4, 3)
+        B = torch.randn(3, 5)
+        result = trntensor.einsum("ij,jk->ik", A, B, dtype=torch.bfloat16)
+        assert result.dtype == torch.bfloat16
+
+    def test_dtype_none_unchanged(self):
+        """dtype=None leaves input dtypes unmodified."""
+        A = torch.randn(4, 3)
+        B = torch.randn(3, 5)
+        result = trntensor.einsum("ij,jk->ik", A, B, dtype=None)
+        assert result.dtype == torch.float32
+
+    def test_dtype_invalid_raises(self):
+        with pytest.raises(ValueError, match="unknown dtype"):
+            trntensor.einsum("ij,jk->ik", torch.randn(4, 3), torch.randn(3, 5), dtype="foo")
+
+    def test_dtype_correctness_bf16(self):
+        """bf16 matmul result is numerically close to fp32 reference."""
+        import numpy as np
+
+        torch.manual_seed(0)
+        A = torch.randn(8, 6)
+        B = torch.randn(6, 10)
+        result_bf16 = trntensor.einsum("ij,jk->ik", A, B, dtype="bf16").float()
+        result_fp32 = trntensor.einsum("ij,jk->ik", A, B)
+        np.testing.assert_allclose(result_bf16.numpy(), result_fp32.numpy(), atol=0.05)

trntensor/einsum.py

@@ -23,13 +23,38 @@
 
 from .plan import ContractionPlan, _parse_subscripts, plan_contraction
 
+# Mapping of user-friendly dtype strings to torch.dtype
+_DTYPE_MAP: dict[str, torch.dtype] = {
+    "bf16": torch.bfloat16,
+    "bfloat16": torch.bfloat16,
+    "fp16": torch.float16,
+    "float16": torch.float16,
+    "f32": torch.float32,
+    "float32": torch.float32,
+    "f64": torch.float64,
+    "float64": torch.float64,
+}
+
+
+def _resolve_dtype(dtype: str | torch.dtype | None) -> torch.dtype | None:
+    """Resolve a dtype argument to a ``torch.dtype`` or ``None``."""
+    if dtype is None:
+        return None
+    if isinstance(dtype, torch.dtype):
+        return dtype
+    try:
+        return _DTYPE_MAP[dtype.lower()]
+    except KeyError:
+        raise ValueError(f"unknown dtype {dtype!r}; valid strings: {sorted(_DTYPE_MAP)}") from None
+
 
 def einsum(
     subscripts: str,
     *operands: torch.Tensor,
     alpha: float = 1.0,
     beta: float = 0.0,
     out: torch.Tensor | None = None,
+    dtype: str | torch.dtype | None = None,
 ) -> torch.Tensor:
     """Einstein summation with contraction planning.
 
@@ -43,20 +68,28 @@ def einsum(
         out: Optional accumulation tensor. When provided the return value is
             ``alpha * contract(operands) + beta * out``. Must have the same
             shape as the contraction result.
+        dtype: Optional compute dtype. When set, all operands are cast to this
+            dtype before contracting and the result is returned in that dtype.
+            Accepts ``torch.dtype`` instances or strings: ``"bf16"``,
+            ``"bfloat16"``, ``"fp16"``, ``"float16"``, ``"f32"``, ``"float32"``.
+            Matches Neuron SDK autocast recommendations for Trainium.
 
     Examples:
         # Matrix multiply
         einsum("ij,jk->ik", A, B)
 
+        # Force bf16 compute (e.g. to hit NKI bf16 matmul kernel)
+        einsum("ij,jk->ik", A, B, dtype="bf16")
+
         # Scaled GEMM: 2*A@B + 0.5*C  (cuTENSOR-style alpha/beta)
         einsum("ij,jk->ik", A, B, alpha=2.0, beta=0.5, out=C)
 
-        # Batched matrix multiply
-        einsum("bij,bjk->bik", A, B)
-
         # DF-MP2 energy contraction
         einsum("iap,jbp->ijab", B, B)
     """
+    target = _resolve_dtype(dtype)
+    if target is not None:
+        operands = tuple(op.to(target) for op in operands)
     plan = plan_contraction(subscripts, *operands)
     result = _execute_contraction(subscripts, operands, plan)
     if alpha != 1.0: