auto is_valid_broadcasting = [](ShapedType operand_ty, ShapedType identity_ty,
                                  ShapedType result_ty) -> bool {
    // Scalar identity is broadcastable to any operand shape, we only need to
    // check that operand has the same shape as a result.
    bool scalar_identity = identity_ty.hasRank() && identity_ty.getRank() == 0;
    if (scalar_identity) return operand_ty == result_ty;

  %0 = "tf.Identity"(%arg0) {_XlaSharding = ""} : (tensor<i32>) -> tensor<i32>
  // CHECK: return %[[OUTPUT]]
  func.return %0 : tensor<i32>
}

// CHECK-LABEL: func @identity_n
// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>, %[[ARG1:.*]]: tensor<f32>)
func.func @identity_n(%arg0: tensor<i32>, %arg1: tensor<f32>) -> (tensor<i32>, tensor<f32>) {
  // CHECK-NOT: "tf.IdentityN"

  void runOnOperation() override;
};

// Returns the first value in the chain of operands, which is not defined by a
// tf.IdentityOp or a tf.ReadVariableOp.
Value SkipIdentityAndReadVariable(Value v) {
  while (auto op = v.getDefiningOp()) {
    if (!isa<TF::IdentityOp, TF::ReadVariableOp>(op)) break;
    v = op->getOperand(0);
  }
  return v;
}

void AnnotateParameterReplicationPass::runOnOperation() {

    DenseFPElementsAttr min_value, max_value;
    if (auto id1 = dyn_cast_or_null<TF::IdentityOp>(min.getDefiningOp())) {
      id1.replaceAllUsesWith(id1.getInput());
      min = tf_op.getMin();
      rewriter.eraseOp(id1);
    }
    if (auto id2 = dyn_cast_or_null<TF::IdentityOp>(max.getDefiningOp())) {
      id2.replaceAllUsesWith(id2.getInput());
      max = tf_op.getMax();

                  AttrType &max_value) const {
    Value min = tf_op.getMin(), max = tf_op.getMax();
    if (auto min_id = min.getDefiningOp<TF::IdentityOp>()) {
      min = min_id.getInput();
    }
    if (auto max_id = max.getDefiningOp<TF::IdentityOp>()) {
      max = max_id.getInput();
    }

    if (!matchPattern(min, m_Constant(&min_value))) {
      return false;
    }

#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {

LogicalResult RemoveIdentity::matchAndRewrite(TF::IdentityOp identity,
                                              PatternRewriter &rewriter) const {
  for (Operation *user : identity->getUsers()) {
    // Replace the op with the input if output is only used by TF ops.

class RankEquals<string rank> : Constraint<CPred<
  "RankEquals($0, " # rank # ")">>;

def IsFusibleWithBias : Constraint<CPred<
  "IsFusibleWithBiasOp($0.getDefiningOp())">>;

// Folds TF IdentityOp with constant input.
def RemoveConstIdentityOp : Pat<
  (TF_IdentityOp (TF_ConstOp $input)),
  (TF_ConstOp $input)>;

// Standardizes the Max and Min ops by moving constant value to rhs. This will

    n = 2 : i32} {
     // CHECK:      %[[IDENTITY_OUT:.*]] = "tf.Identity"(%[[RI_0]])
     // CHECK:      %[[RESOURCE_OUT:.*]] = "tf_device.launch"()
     // CHECK-SAME: TPU_REPLICATED_CORE_0
     // CHECK-NEXT:   %[[READ_OUT:.*]] = "tf.ReadVariableOp"(%[[IDENTITY_OUT]])
     // CHECK-NEXT:   tf_device.return %[[READ_OUT]]

  // isn't a composite op. The following ops are explicitly skipped here because
  // their "no-op" expansion is known to cause problems in some cases.
  static const char* kOpsToSkip[] = {
      "IdentityOp",
      "NoOp",              // b/174596063
      "OptionalHasValue",  // b/173136483
      "OptionalGetValue",  // b/173136483
      "VarHandleOp",       // b/176819198
  };
  for (const char* skip : kOpsToSkip) {

# CHECK: return %[[IDENTITY]]#1, %[[IDENTITY]]#0


def Test():

  x = tf.constant(1.0, shape=(5, 3))
  y = tf.constant(1.0, shape=(3, 5))

  s = tf.matmul(x, y)
  t = tf.matmul(y, x)
  [t, s] = array_ops.identity_n([t, s])

  tensor_info_x = tf.compat.v1.saved_model.utils.build_tensor_info(x)
  tensor_info_y = tf.compat.v1.saved_model.utils.build_tensor_info(y)
  tensor_info_s = tf.compat.v1.saved_model.utils.build_tensor_info(s)