einsum support

arogozhnikov · Dec 9, 2023 · 194dad8 · 194dad8
1 parent 970ed5a
commit 194dad8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/einops/_backends.py b/einops/_backends.py
@@ -708,4 +708,7 @@ def concat(self, tensors, axis: int):
         return tensors[0].cat(tensors[1:], axis) if len(tensors) > 1 else tensors[0]
 
     def is_float_type(self, x):
-        return self.tinygrad.dtypes.is_float(x.dtype)
+        return self.tinygrad.dtypes.is_float(x.dtype)
+
+    def einsum(self, pattern, *x):
+        return self.tinygrad.Tensor.einsum(pattern, *x)