Add TF2XLA registration for Einsum op, so that it uses XlaEinsum implementation. (Try #2)

This is needed because we want to switch to the fused EinsumOp in TF Classic when from the python implementation of tf.einsum. After this, the registration would enable compilation of the new TF Graphs in XLA.

Note: It will still emit XlaEinsum and not Einsum when enclosing context is TPU.
PiperOrigin-RevId: 268083168

Author: bloops

tensorflow/compiler/tf2xla/kernels/einsum_op.cc

@@ -48,6 +48,7 @@ class EinsumOp : public XlaOpKernel {
 };
 
 REGISTER_XLA_OP(Name("XlaEinsum").TypeConstraint("T", kEinsumTypes), EinsumOp);
+REGISTER_XLA_OP(Name("Einsum").TypeConstraint("T", kEinsumTypes), EinsumOp);
 
 }  // namespace
 }  // namespace tensorflow

tensorflow/python/kernel_tests/einsum_op_test.py

@@ -50,6 +50,7 @@ def _check(self, s, *input_shapes, **kwargs):
     b = self.evaluate(gen_linalg_ops.einsum(input_tensors, s))
     self.assertAllClose(a, b, atol=1e-4, rtol=1e-4)
 
+  @test_util.disable_xla('b/131919749')
   def testUnary(self):
     self._check('->', ())
     self._check('aa->', (3, 3))
@@ -66,7 +67,10 @@ def testUnary(self):
     self._check('aabcc->ac', (3, 3, 5, 4, 4))
     self._check('aabcd->ad', (3, 3, 5, 4, 4))
 
+  @test_util.disable_xla('b/131919749')
   def testUnaryEllipsis(self):
+    # Unary cases with ellipsis.
+    # Edge cases.
     self._check('...->...', ())
     self._check('...->', ())
     self._check('->...', ())
@@ -78,31 +82,44 @@ def testUnaryEllipsis(self):
     self._check('a...a->a...', (2, 1, 2))
     self._check('a...a->a...', (2, 3, 4, 5, 2))
 
+    # Regular cases.
     self._check('...ijk->...ki', (3, 4, 5))
     self._check('...ijk->...ki', (1, 3, 4, 5))
     self._check('...ijk->...ki', (2, 2, 3, 4, 5))
 
     # Repeated indices.
     self._check('i...ii->...i', (3, 2, 3, 3))
 
-  def testBinary(self):
+  def testBinarySimple(self):
+    # Binary cases in XLA mode must have either (a) each index appearing exactly
+    # once in both the inputs (batch or contraction index), or (b) appearing
+    # exactly once in an input and in the output (free index).
     self._check(',->', (), ())
     self._check('a,a->', (3,), (3,))
     self._check('a,a->a', (3,), (3,))
-    self._check('ba,b->', (3, 2), (3,))
     self._check('ab,b->a', (3, 4), (4,))
     self._check('ab,ab->', (3, 4), (3, 4))
+    self._check('ab,bc->ac', (3, 4), (4, 5))
     self._check('nij,jk->nik', (5, 2, 3), (3, 4))
     self._check('abc,bad->abcd', (1, 2, 3), (2, 1, 4))
+    # Based on https://github.com/google/jax/issues/37#issuecomment-448572187
+    self._check('sa,shb->shab', (2, 1), (2, 3, 4))
+
+  @test_util.disable_xla('b/131919749')
+  def testReducedIndices(self):
+    self._check('ba,b->', (3, 2), (3,))
+    self._check('ab,ab->', (3, 4), (3, 4))
+
+  @test_util.disable_xla('b/131919749')
+  def testRepeatedIndices(self):
     # Repeated indices.
     self._check('ijj,k->ik', (2, 3, 3), (4,))
     self._check('aba,a->b', (3, 4, 3), (3,))
     # From https://github.com/dask/dask/pull/3412#discussion_r182413444
     self._check('aab,bc->ac', (2, 2, 3), (3, 4))
     self._check('aab,bcc->ac', (2, 2, 3), (3, 4, 4))
-    # Based on https://github.com/google/jax/issues/37#issuecomment-448572187
-    self._check('sa,shb->shab', (2, 1), (2, 3, 4))
 
+  @test_util.disable_xla('b/131919749')
   def testBroadcasting(self):
     # Batch matmul without broadcasting.
     self._check('...ij,...jk->...ik', (5, 1, 2, 3), (5, 1, 3, 4))
@@ -113,14 +130,17 @@ def testBroadcasting(self):
     self._check('...ij,...jk->...ik', (2, 3), (5, 3, 5))
     self._check('...ij,...jk->...ik', (3, 1, 2, 3), (1, 1, 7, 3, 5))
     self._check('i...j,j...k->...ik', (2, 1, 3, 1, 3), (3, 1, 7, 5))
+    # Following 2 from # https://stackoverflow.com/a/19203475/1611416
+    self._check('...abc,...abcd->...d', (1, 1, 2, 3, 4), (5, 2, 3, 4, 6))
+    self._check('ab...,b->ab...', (2, 3, 1, 1, 5), (3,))
+
+  @test_util.disable_xla('b/131919749')
+  def testBroadcastingWithRepeatedIndices(self):
     # Broadcasting with repeated indices.
     self._check('ij,jk...k->i...', (3, 2), (2, 4, 1, 4))
     self._check('ij,jk...k->...i', (3, 2), (2, 4, 5, 4))
     self._check('ijj,jk...k->i...', (3, 2, 2), (2, 4, 1, 4))
     self._check('i...jj,jk...k->i...', (3, 3, 1, 2, 2), (2, 4, 1, 5, 4))
-    # Following 2 from # https://stackoverflow.com/a/19203475/1611416
-    self._check('...abc,...abcd->...d', (1, 1, 2, 3, 4), (5, 2, 3, 4, 6))
-    self._check('ab...,b->ab...', (2, 3, 1, 1, 5), (3,))
 
   def testDtypes(self):
     bfloat16 = dtypes.bfloat16.as_numpy_dtype
@@ -153,6 +173,7 @@ def check(dtype):
     ]:
       check(dtype)
 
+  @test_util.disable_xla('b/131919749')
   @test_util.run_in_graph_and_eager_modes
   def testInvalid(self):
     r = np.random.RandomState(0)
@@ -178,6 +199,7 @@ def testInvalid(self):
       with self.assertRaises((ValueError, errors.InvalidArgumentError)):
         _ = self.evaluate(gen_linalg_ops.einsum(placeholders, args[0]))
 
+  @test_util.disable_xla('b/131919749')
   @test_util.run_in_graph_and_eager_modes
   def testPlaceholder(self):
 
@@ -202,9 +224,11 @@ def check(equation, *input_and_placeholder_shapes):
           ((4, 3), (None, 3)))
     check('...ij,...jk->...ik', ((3, 1, 2, 3), None), ((1, 7, 3, 4), None))
 
+  @test_util.disable_xla('b/131919749')
   def testOutputRepeatedLabels(self):
-    # This is the reverse operation of repeated input labels, to be used for
-    # computing symbolic gradients of einsum.
+    # This is the reverse operation of generalized traces, to be used for
+    # computing symbolic gradients of einsum. Note: this operation is not
+    # supported by np.einsum as it's only required for gradients.
     r = np.random.RandomState(0)
     a = r.randn(2, 2)
     s = 'a->aa'