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.

def IsPrecisionEmpty :
  Constraint<CPred<"IsPrecisionEmpty($0)">>;

// Creates Einsum Op from XlaDotV2 Op by generating equation.
def CreateEinsumOpFromXlaDotV2Op : NativeCodeCall<
  "CreateEinsumOpFromXlaDotV2Op($_builder, $_loc, $0...)">;

// Convert XlaDotV2 Op to Einsum Op with above two functions.
def ConvertXlaDotV2OpToEinsumOp : Pat<

# REQUIREMENTS_FILE=jax.requirements.txt


setuptools
wheel
cloudpickle
colorama>=0.4.4
matplotlib
pillow>=9.1.0
rich
absl-py
portpicker
six
opt-einsum
auditwheel
typing_extensions
importlib_metadata>=4.6
numpy==1.26.0;python_version=="3.12"
numpy==1.23.4;python_version=="3.11"
numpy==1.22.4;python_version<"3.11"
scipy==1.11.2;python_version=="3.12"

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/core/util/matmul_bcast.h"

namespace mlir {
namespace TF {

namespace {

// Replace TF BatchMatMul by TF Einsum op
template <typename BatchMatMulOpType>
class ConvertTFBatchMatMulToEinsumOp
    : public OpRewritePattern<BatchMatMulOpType> {
  using OpRewritePattern<BatchMatMulOpType>::OpRewritePattern;

  CreateModule(kNonStaticFailure);

  auto result = Run();

  EXPECT_TRUE(result.succeeded());
  EXPECT_EQ(static_error.Delta("mhlo.einsum"), 0);
  EXPECT_EQ(skipped.Delta("mhlo.einsum"), 1);
}

TEST_F(VerifyTfxlaLegalizationTest, SkipsNonStaticInputsWithBounds) {
  // Using a string constant here instead of testdata to make this compatible
  // with open source.

bool IsInLiftedFunc(Operation* op);

// Checks if the op is inside a StableHLO op with region.
// If the given op pointer is a nullptr, returns false.
bool IsInStableHloOpRegion(Operation* op);

// Checks if a given einsum op is supported for XlaDotV2 quantization.
bool IsEinsumSupportedByXlaDotV2(StringAttr equation_attr);

// Gets the quantization method from `op`. It is retrieved from the
// `kQuantizationMethodAttr` string attribute. Returns

  %0 = mhlo.constant dense<(1.000000e+00,-1.000000e+00)> : tensor<128x32x4xcomplex<f32>>
  %1 = mhlo.constant dense<(1.000000e+00,1.000000e+00)> : tensor<8x64x128xcomplex<f32>>
  %2 = "mhlo.einsum"(%1, %0) <{einsum_config = "abc,cde->abde"}> : (tensor<8x64x128xcomplex<f32>>, tensor<128x32x4xcomplex<f32>>) -> tensor<8x64x32x4xcomplex<f32>>
  return %2 : tensor<8x64x32x4xcomplex<f32>>
}

// -----