%16 = stablehlo.broadcast_in_dim %6, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x4xi32>
    %17 = stablehlo.subtract %15, %16 : tensor<8x16x4xi32>  // q2 - z2
    // Corresponds to einsum expression: b i j, b j d -> b i d
    %18 = stablehlo.dot_general %14, %17, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x16xi32>, tensor<8x16x4xi32>) -> tensor<8x16x4xi32>

SmallVector<Value> AppendToVector(const ArrayRef<Value> arguments,
                                  Value append) {
  SmallVector<Value> ret(arguments);
  ret.push_back(append);
  return ret;
}

// Check if the given einsum equation is supported by XlaDotV2.
// Conditions:
// 1. Two inputs & one output.
// 2. No ... in the equation.
// 3. Batch dimensions should be the same, or only the left equation should have

          quantize_func_count++;
        } else if (func_name.starts_with("dequantize_i")) {
          dequantize_func_count++;
        }
      } else if (auto einsum = llvm::isa<TF::EinsumOp>(op)) {
        if (IsInCompsiteFunction(op)) return;
        // Leftover Einsum ops are always non-quantized.
        auto op_name = op->getName().stripDialect();
        func_count_map[op_name].num_float++;
      }
    });

    MLIR pipeline in the future.
  }];
  let constructor = "TF::CreateGpuOpFusionPass()";
}

def BatchMatMulToEinsumPass : Pass<"tf-batch-matmul-to-tf-einsum", "mlir::func::FuncOp"> {
  let summary = "Replace TF BatchMatMul op by TF Einsum op.";
  let constructor = "TF::CreateBatchMatMulToEinsumPass()";
}

TensorFlow graph export by making transformation such as replacing or
removing MLIR or XLA specific attributes that are not legal in TensorFlow
graph.
### `-tf-batch-matmul-to-tf-einsum`

_Replace TF BatchMatMul op by TF Einsum op._

### `-tf-broadcast-fold`

_Fold explicit broadcasts into the following operations if they support implicit broadcasting on their operand._

// CHECK-NEXT: return %[[BIASADD]] : tensor<?x?x?x3xf32>

// -----

func.func @lower_einsum(%arg0: tensor<3x4x5xf32>, %arg1: tensor<3x5x6xf32>) -> tensor<3x4x6xf32> {
  %0 = "tf.Einsum"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", equation = "ijk,ikm->ijm"}: (tensor<3x4x5xf32>, tensor<3x5x6xf32>) -> tensor<3x4x6xf32>
  func.return %0 : tensor<3x4x6xf32>
}
// CHECK-LABEL: lower_einsum

}

//===----------------------------------------------------------------------===//
// Einsum.
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @einsum
func.func @einsum(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> {
  // CHECK:  mhlo.einsum
  %0 = "tf.Einsum"(%arg0, %arg1) {equation = "ab,bc->ac"} : (tensor<2x3xf32>, tensor<3x4xf32>) -> tensor<2x4xf32>

#include "tensorflow/compiler/mlir/lite/utils/utils.h"
#include "tensorflow/compiler/mlir/lite/utils/validators.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/einsum.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
#include "tensorflow/compiler/mlir/tensorflow/transforms/unroll_batch_matmul.h"
#include "tensorflow/compiler/mlir/tensorflow/utils/dynamic_shape_utils.h"

is raised.

@compatibility(numpy)
Similar to [`numpy.einsum`](https://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html).

Comparison with `numpy.einsum`:

 * This Op only supports unary and binary forms of `numpy.einsum`.
 * This Op does not support implicit form. (i.e. equations without `->`).

      "Dequantize",
      "Diag",
      "DynamicInfeedEnqueueTupleOp",
      "DynamicInfeedDequeueTupleOp",
      "DynamicStitch",
      "DynamicPartition",
      "Einsum",
      "EmptyTensorList",
      "EnsureShape",
      "ExtractImagePatches",
      "Igamma",
      "IgammaGradA",
      "RandomGammaGrad",
      "Igammac",
      "FFT",
      "FFT2D",