Add abstract method

AI-Hypercomputer · Feb 11, 2025 · c67c7ca · c67c7ca
1 parent adab3f6
commit c67c7ca
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/MaxText/layers/quantizations.py b/MaxText/layers/quantizations.py
@@ -63,6 +63,10 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()):
     """Placeholder for dot_general implementation in subclasses."""
     pass
 
+  def einsum(self, dtype: DType = jnp.float32):
+    """Placeholder for einsum implementation in subclasses."""
+    pass
+
 
 def _tiling_fn(lhs, rhs, dimension_numbers, tile_size):
   del lhs, rhs