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

  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.

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

// -----

      }

      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) {

          function_name.contains("and_bias")) {
        spec->biases_params[2] = {{0, 1},
                                  quant::GetUniformQuantizedTypeForBias};
      }
    } else if (function_name.contains("einsum")) {
      spec->coeff_op_quant_dim[1] = -1;
      if (function_name.contains("with_bias")) {
        spec->biases_params[2] = {{0, 1},
                                  quant::GetUniformQuantizedTypeForBias};
      }

#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