]
      )
      def einsum_without_kernel(
          self, x: core.Tensor, y: core.Tensor
      ) -> Mapping[str, core.Tensor]:
        return self._einsum(x, y)

      def _einsum(self, x, y):

        out = tensorflow.einsum(equation, x, y)
        if self._bias is not None:
          out = nn_ops.bias_add(out, self._bias)
        return {'output': out}

    model = EinsumModel()

              ),
          ]
      )
      def einsum_without_kernel(
          self, x: core.Tensor, y: core.Tensor
      ) -> Mapping[str, core.Tensor]:
        return self._einsum(x, y)

      def _einsum(self, x, y):
        if is_qat_model:
          x = array_ops.fake_quant_with_min_max_vars(
              x,
              min=ops.convert_to_tensor(self._min[0]),

// RUN: tf-opt -split-input-file -verify-diagnostics -tf-einsum %s | FileCheck %s

func.func @unary_einsum_reduce_sum_transpose(%arg0: tensor<3x4x5x6xf32>) -> tensor<3x5x4xf32> {
  %0 = "tf.Einsum"(%arg0) {T = "tfdtype$DT_FLOAT", equation = "...gse->...sg"}: (tensor<3x4x5x6xf32>) -> tensor<3x5x4xf32>
  func.return %0 : tensor<3x5x4xf32>
  // CHECK-LABEL: unary_einsum_reduce_sum_transpose
  // CHECK-DAG: %[[cst:.*]] = arith.constant dense<3> : tensor<1xi32>

}

struct EinsumDimensionNumbers {
  // Each field contains the list of dimensions appearing only in the specifed
  // arguments of the einsum op with natural ordering. For example `rhs_out`
  // contains the dimensions appearing in the RHS and the OUTPUT of the einsum
  // but not in the LHS.
  std::vector<int64_t> lhs;
  std::vector<int64_t> rhs;
  std::vector<std::tuple<int64_t, int64_t>> lhs_rhs;

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This pass identifies patterns for certain Einsum Ops and replaces them
// with other equivalent TF Ops.

#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_EINSUM_H_
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_EINSUM_H_

#include <cstdint>

  // Claim: For the einsum operation z = einsum("{eq_x},{eq_y}->{eq_z}", x, y),
  //   where the equation involves only Tensor contractions, generalized traces
  //   and transposes, the input gradients are given by the vector-jacobian
  //   products (VJPs):
  //
  //     grad_wrt_x = einsum("{eq_y},{eq_z}->{eq_x}", y, grad_wrt_z)
  //     grad_wrt_y = einsum("{eq_x},{eq_z}->{eq_y}", x, grad_wrt_z}
  //

  auto z = Einsum(scope_, {x}, "ii->");
  TensorShape z_shape({});
  RunTest({x}, {x_shape}, {z}, {z_shape});
}

TEST_F(LinalgGradTest, Einsum_TraceBroadcast) {
  TensorShape x_shape({4, 3, 3});
  Output x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
  // Note: In Python this could just be "ii" becuase tf.einsum normalizes the
  // equation, but c++ doesn't do that.
  auto z = Einsum(scope_, {x}, "...ii->...");

// RUN: tf-opt %s -tf-batch-matmul-to-tf-einsum | FileCheck %s

func.func @test_batch_matmul_to_einsum(%arg0: tensor<1x2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<1x2x4xf32> {
  // CHECK-LABEL: test_batch_matmul_to_einsum
  // CHECK: "tf.Einsum"(%arg0, %arg1) <{equation = "...mk,...kn->...mn"}> : (tensor<1x2x3xf32>, tensor<3x4xf32>) -> tensor<1x2x4xf32>

  func.return %0 : tensor<?x2x4x5xf32>
}

// CHECK: func @xla_dot_v2
// CHECK: %[[einsum:.*]] = "tf.Einsum"(%arg0, %arg1) <{equation = "abc,cde->abde"}> : (tensor<?x2x3xf32>, tensor<3x4x5xf32>) -> tensor<?x2x4x5xf32>
// CHECK: return %[[einsum]] : tensor<?x2x4x5xf32>

// -----

// dimension_numbers: {
//   offset_dims: 0
//   collapsed_slice_dims: 1
//   start_index_map: 1

                    "\""# func_name #"\", $0...)", returns>;

// Add the second argument to the first argument, which is expected to be an
// argument list.
// bias(einsum(inputs), bias) --> einsum_with_bias(AppendToVector(inputs, bias))
// Since inputs is a vector in case of einsum, we cannot use ArgumentList here.
def AppendToVector : NativeCodeCall<"AppendToVector($0, $1)">;

// The list of arguments of the composite function.