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

#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:

    # 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

    # 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

    # 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

    # 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

        "decompose_reduce_dataset.cc",
        "decompose_resource_ops_pass.cc",
        "device_attribute_to_launch.cc",
        "device_index_selector.cc",
        "drop_while_shape_invariant.cc",
        "einsum.cc",
        "executor_island_coarsening.cc",
        "executor_tpuv1_inline_tpu_island.cc",
        "executor_tpuv1_island_coarsening.cc",
        "executor_tpuv1_outline_tpu_island.cc",

std::unique_ptr<OperationPass<func::FuncOp>> CreateUnrollBatchMatMulPassPass();

// Optional pass which will map TF BatchMatMul to TF Einsum
std::unique_ptr<OperationPass<func::FuncOp>> CreateBatchMatMulToEinsumPass();

// Pass that transform Einsum to other TF Ops for the supported variants.
std::unique_ptr<OperationPass<func::FuncOp>> CreateTransformEinsumPass();

// Optimizes Tensorflow graph.

    %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