Value input_rhs = bmm_op.getY();

    Value output_lhs =
        bmm_op.getAdjX() ? create_z_x_transpose_op(input_lhs) : input_lhs;

    // The rhs need to be transposed if adj_y == false AND this matmul will be
    // legalized to tfl.fully_connected
    Value output_rhs =
        !bmm_op.getAdjY() ? create_z_x_transpose_op(input_rhs) : input_rhs;

    Type output_type = bmm_op.getResult().getType();

    return failure();
  }

  // Replace the last 2 dimensions of LHS and RHS if necessary.
  // The actual transpose is done by MatMulOp.
  if (op.getAdjX()) {
    std::swap(lhs_shape[lhs_dims - 1], lhs_shape[lhs_dims - 2]);
  }
  if (op.getAdjY()) {
    std::swap(rhs_shape[rhs_dims - 1], rhs_shape[rhs_dims - 2]);
  }

  const int64_t rows = lhs_shape[lhs_dims - 2];

        }

        const int x_row =
            matmul_op.getAdjX() ? shape_x.back() : *(shape_x.rbegin() + 1);
        const int x_col =
            !matmul_op.getAdjX() ? shape_x.back() : *(shape_x.rbegin() + 1);

        const int y_row =
            matmul_op.getAdjY() ? shape_y.back() : *(shape_y.rbegin() + 1);
        const int y_col =
            !matmul_op.getAdjY() ? shape_y.back() : *(shape_y.rbegin() + 1);

    if (dims_a < 2 || dims_b < 2) {
      return failure();
    }

    // einsum equation for batchmatmul
    std::string equation("...mk,...kn->...mn");
    if (op.getAdjX()) std::swap(equation[3], equation[4]);
    if (op.getAdjY()) std::swap(equation[6 + 3], equation[6 + 4]);

    rewriter.replaceOpWithNewOp<TF::EinsumOp>(
        op, op.getType(),
        /*inputs=*/ValueRange({input_lhs, input_rhs}),

  Value input_rhs = get_real_input_value(op.getY());

  Value legalized_lhs =
      op.getAdjX() ? create_z_x_transpose_op(input_lhs) : input_lhs;

  // The rhs need to be transposed if adj_y == false AND this matmul will be
  // legalized to tfl.fully_connected
  Value legalized_rhs =
      !op.getAdjY() ? create_z_x_transpose_op(input_rhs) : input_rhs;

  Type output_type = op.getResult().getType();

    auto new_reshape = rewriter.create<ReshapeOp>(
        op.getLoc(), transpose_op.getInput(), shape_constant);
    rewriter.replaceOpWithNewOp<BatchMatMulOp>(
        op, op.getType(), op.getX(), new_reshape, op.getAdjX(), !op.getAdjY(),
        op.getAsymmetricQuantizeInputsAttr());
    return success();
  }

 private:
  // Checks that tensor `a` has shape of [M, N, ...] and `b` has [M * N, ...].

    if (!lhs_type || !rhs_type) return failure();
    if (mlir::isa<ComplexType>(lhs_type.getElementType()) && op.getAdjX()) {
      lhs = rewriter.create<TF::ConjOp>(op.getLoc(), lhs_type, lhs);
    }
    if (mlir::isa<ComplexType>(rhs_type.getElementType()) && op.getAdjY()) {
      rhs = rewriter.create<TF::ConjOp>(op.getLoc(), rhs_type, rhs);
    }

    // Broadcast both operands.

  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

  int64_t expected_out_row_dim = op.getAdjX() ? x_col_dim : x_row_dim;
  int64_t expected_out_col_dim = op.getAdjY() ? y_row_dim : y_col_dim;

  if (expected_out_row_dim != ShapedType::kDynamic &&
      out_row_dim != ShapedType::kDynamic &&
      out_row_dim != expected_out_row_dim)

  int64_t out_row_dim = output_shape[output_shape.size() - 2];
  int64_t out_col_dim = output_shape[output_shape.size() - 1];

  int64_t expected_out_row_dim = op.getAdjX() ? x_col_dim : x_row_dim;
  int64_t expected_out_col_dim = op.getAdjY() ? y_row_dim : y_col_dim;

  if (expected_out_row_dim != ShapedType::kDynamic &&
      out_row_dim != ShapedType::kDynamic &&
      out_row_dim != expected_out_row_dim)