]
      )
      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()

// 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;

  // 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}
  //

  [], (addBenefit 10)>;

// Converts inlined Einsum pattern to TF XlaDotV2 op.
def ConvertTFEinsumToXLADotV2Op : Pat<
  (TF_EinsumOp:$einsum
    $args, $equation),
  (CreateXlaDotV2OpFromTfEinsumOp
    $equation, $args, $einsum),
  [(IsInt32ElementType $einsum),
   // Constraint to check:
   // 1. The einsum has two inputs and one output.
   // 2. The einsum is not created by the convert function itself.

  if (!value_type.hasRank()) return false;
  if (!value_type.getElementType().isInteger(integer_width)) return false;

  return true;
}

// Constraint to check:
// 1. The einsum has two inputs and one output.
// 2. The einsum is not created by the convert function itself.
// 3. Both inputs are int32 tensor.
// 4. Both inputs have the graph ancestor of either const-(sub), or cast-sub.

      }

      if (!is_weight_constant) {
        if (!function_name.contains("matmul") &&
            !function_name.contains("einsum")) {
          return absl::InternalError(
              "Non-constant weights are not supported at the moment,"
              " except matmul and einsum.");
        } else if (!quant_options_.enable_two_input_tensors() &&
                   !is_unitwise_quantization_enabled) {

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/einsum.h"

namespace mlir {
namespace quant {
namespace {

using ::tensorflow::quantization::OpSet;

class PrepareLiftingPass
    : public PassWrapper<PrepareLiftingPass, OperationPass<func::FuncOp>> {
 public:

    func.return %5 : tensor<*xi32>
  }

   // Einsum with int32 accumulation.
  func.func private @internal_einsum_fn(
                         %input : tensor<*xi8>, %weight : tensor<*xi8>,
                         %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>,
                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xi32> {

    # via
    #   -r requirements.in
    #   h5py
    #   jax
    #   keras-nightly
    #   ml-dtypes
    #   opt-einsum
    #   scipy
    #   tb-nightly
opt-einsum==3.3.0 \
    --hash=sha256:2455e59e3947d3c275477df7f5205b30635e266fe6dc300e3d9f9646bfcea147 \
    --hash=sha256:59f6475f77bbc37dcf7cd748519c0ec60722e91e63ca114e68821c0c54a46549
    # via