}

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, arith::ArithDialect>();
  }

  void runOnOperation() override;
};

// Generate an einsum equation from the given DotDimensionNumber.
std::string CreateEinsumEquation(
    const xla::DotDimensionNumbers& dot_dimension_numbers, const int lhs_rank,
    const int rhs_rank) {
  // Prepare necessary indices.

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