}
}

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

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

          callee.getArgument(use.getOperandNumber()), module, infer_from_op);
      if (type_from_callee.has_value()) return type_from_callee;
    } else if (llvm::isa<TF::IdentityOp, TF::IdentityNOp>(use.getOwner())) {
      auto type_from_alias = GetElementTypeFromAccess(
          use.getOwner()->getResult(use.getOperandNumber()), module,
          infer_from_op);

                                   val_bcast.getType(), block_arg))) {
    return failure();
  }

  OpBuilder before_cluster_builder(cluster);
  IdentityOp assigned_id = before_cluster_builder.create<IdentityOp>(
      val_bcast.getLoc(), block_arg.getType(), block_arg);
  assigned_id->setAttr(kICIWeightDistributionMlirBridgeMarker,
                       before_cluster_builder.getBoolAttr(true));

    %identity_out, %control_identity = tf_executor.island wraps "tf.Identity"(%replicated_out) {device = ""} : (tensor<4x4xf32>) -> tensor<4x4xf32>
    %partitioned_out:2, %partitioned_control = tf_executor.island wraps "tf.TPUPartitionedOutput"(%identity_out) {partition_dim = 0 : i64} : (tensor<4x4xf32>) -> (tensor<2x4xf32>, tensor<2x4xf32>)
    tf_executor.fetch

    const Block& func_block = func.front();
    // Walk all operations and nested regions to find accessed gradient sources
    // for function arguments.
    func.walk([&](Operation* op) {
      if (llvm::isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
        op->replaceAllUsesWith(op->getOperands());
        return;
      }
      if (auto grad = llvm::dyn_cast<TF::TensorArrayGradV3Op>(op)) {

        decomposed_partitioned_call_callees) {
  for (auto& op : llvm::make_early_inc_range(block->getOperations())) {
    // TODO(yuanzx): Add a pass to remove identities in device computation.
    if (llvm::isa<TF::IdentityOp, TF::IdentityNOp, TF::StopGradientOp>(&op)) {
      op.replaceAllUsesWith(op.getOperands());
      op.erase();
    } else if (auto list = llvm::dyn_cast<TF::EmptyTensorListOp>(&op)) {

// such nodes to carry information.
void RemovePassthroughOp(Block &block) {
  for (auto &op : llvm::make_early_inc_range(block)) {
    if (isa<TF::IdentityOp, TF::IdentityNOp>(op) || IsCastOfResource(op)) {
      op.replaceAllUsesWith(op.getOperands());
      op.erase();
    }
  }
}

using LocalVarOp = std::variant<TF::VarHandleOp, TF::MlirLocalVarOp>;