// shape information which is necessary for correct pattern matching in this
// pass.
struct RemoveIdentity : public OpRewritePattern<TF::IdentityOp> {
  using OpRewritePattern<TF::IdentityOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TF::IdentityOp identity,
                                PatternRewriter &rewriter) const override;
};

}  // namespace quant
}  // namespace mlir

// CHECK-DAG:   %[[IDENTITY_0:.*]] = "tf.Identity"(%[[PARTITIONEDCALL_0]])
// CHECK-DAG:   %[[IDENTITY_1:.*]] = "tf.Identity"(%[[PARTITIONEDCALL_1]])
// CHECK:   return %[[IDENTITY_0]], %[[IDENTITY_1]] : tensor<1xf32>, tensor<1xf32>
// CHECK: }
}

// -----

  }
}

// Checks if an operation can have TPU devices propagated through.
bool IsSupportedOpToSetDevice(Operation& op) {
  return IsSupportedExecutorOp(op) ||
         isa<TF::IdentityOp, TF::IdentityNOp, TF::ShapeOp>(op);
}

// Finds nonconflicting TPU device for an operation from its operands. If an
// operand has no device or a non TPU device, or if there are conflicting

  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;

      }
    }
    node_def {
      name: "Identity_1"
      op: "Identity"
      input: "cond_true1"
      attr {
        key: "T"
        value {
          type: DT_INT32
        }
      }
    }
    ret {
      key: "cond_true2"
      value: "Identity:output:0"
    }
    ret {
      key: "cond_true3"
      value: "Identity_1:output:0"
    }
  }
  function {

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

								}
							}
							cacheKey := utils.ToStringKey(relPrimaryValues...)
							if len(relPrimaryValues) != len(rel.FieldSchema.PrimaryFields) || !identityMap[cacheKey] {
								if cacheKey != "" { // has primary fields
									identityMap[cacheKey] = true
								}

								distinctElems = reflect.Append(distinctElems, rv)
							}
						}
					}

					if elems.Len() > 0 {

      // in the Island body.
      if (value == island.getControl()) break;
      value = island.GetYield().getOperand(res_index);
    } else if (isa<IdentityNOp, IdentityOp>(op)) {
      value = op->getOperand(res_index);
    } else if (auto call = dyn_cast<CallOpInterface>(op)) {
      func::FuncOp func = dyn_cast<func::FuncOp>(

			}

			if results, ok := joinIdentityMap[utils.ToStringKey(fieldValues...)]; ok {
				joinKey := utils.ToStringKey(joinFieldValues...)
				identityMap[joinKey] = append(identityMap[joinKey], results...)
			}
		}

		_, foreignValues = schema.GetIdentityFieldValuesMap(tx.Statement.Context, joinResults, joinRelForeignFields)
	} else {
		for _, ref := range rel.References {