op->emitOpError()
              << "is missing attribute '" << kTfBackendConfigAttrName << "'";
          return WalkResult::interrupt();
        }

        auto called_index_attr = mlir::dyn_cast_or_null<IntegerAttr>(
            backend_config.get(kCalledIndexAttrName));
        if (!called_index_attr) {
          op->emitOpError()
              << "is missing attribute '" << kCalledIndexAttrName << "'";

    Type elem_type = getType().cast<ShapedType>().getElementType();
    if (elem_type.isSignlessInteger()) {
      auto start_attr = start_tensor.getValues<IntegerAttr>()[0];
      auto limit_attr = limit_tensor.getValues<IntegerAttr>()[0];
      auto delta_attr = delta_tensor.getValues<IntegerAttr>()[0];
      const int num_elements = GetLengthOfRange(
          start_attr.getInt(), limit_attr.getInt(), delta_attr.getInt());

    use_spmd = use_spmd_attr.getValue();

  if (!use_spmd) return success();

  auto num_cores_per_replica_attr =
      cluster_func->getAttrOfType<IntegerAttr>(kNumCoresPerReplicaAttr);

  // Wrap the ClusterFunc with a ParallelExecute if it does not already exist.
  OpBuilder builder(cluster_func);
  tf_device::ParallelExecuteOp parallel_execute =

    if (!num_replicas || !num_replicas.isa<mlir::IntegerAttr>())
      return cluster.emitError()
             << "requires '" << kNumReplicasAttr << "' int attribute";

    int num_cores_per_replica = 1;
    auto num_cores_per_replica_attr = mlir::dyn_cast_or_null<mlir::IntegerAttr>(
        cluster_metadata->getSecond().get(kNumCoresPerReplicaAttr));
    if (num_cores_per_replica_attr)

      result = ConstFoldBroadcastInDim<FloatAttr>(result_type, const_val,
                                                  bcast_dims);
    } else if (mlir::isa<IntegerType>(elem_type)) {
      result = ConstFoldBroadcastInDim<IntegerAttr>(result_type, const_val,
                                                    bcast_dims);
    } else {
      return rewriter.notifyMatchFailure(bcast_op, "Unsupported element type.");
    }
    Value new_const_op =

                               output_arg.name(), "'");
      if (!mlir::isa<IntegerAttr>(repeats_attr))
        return InvalidArgument("Attribute '", output_arg.number_attr(),
                               "' required for output list '",
                               output_arg.name(), "' isn't an integer");
      int64_t repeats = mlir::cast<IntegerAttr>(repeats_attr).getInt();

      if (!output_arg.type_attr().empty()) {

          mlir::Attribute rand_val;
          mlir::Type element_type = attr.getShapedType().getElementType();
          if (mlir::isa<mlir::IntegerType>(element_type)) {
            rand_val = mlir::IntegerAttr::get(element_type, std::rand());
          } else if (element_type.isF16() || element_type.isF32() ||
                     element_type.isF64()) {
            rand_val = mlir::FloatAttr::get(element_type,

  llvm::SmallSet<int, 8> read_array;
  for (const Attribute& attr : op.getDeviceVarReadsIndices()) {
    read_array.insert(mlir::cast<IntegerAttr>(attr).getInt());
  }
  llvm::SmallSet<int, 8> update_array;
  for (const Attribute& attr : op.getDeviceVarUpdatesIndices()) {
    update_array.insert(mlir::cast<IntegerAttr>(attr).getInt());
  }

  for (auto& arg : op->getOpOperands()) {
    Value v = arg.get();

  unsigned num_arguments = main_fn.getNumArguments();
  for (unsigned arg_index = 0; arg_index < num_arguments; ++arg_index) {
    if (auto aliasing_output = main_fn.getArgAttrOfType<mlir::IntegerAttr>(
            arg_index, kAliasingAttr))
      output_to_input_alias[aliasing_output.getInt()] = arg_index;
  }

  if (output_to_input_alias.empty()) return absl::OkStatus();

  using OpRewritePattern<GetElementOp>::OpRewritePattern;

 public:
  LogicalResult matchAndRewrite(GetElementOp ge_op,
                                PatternRewriter &rewriter) const override {
    IntegerAttr index;
    if (!matchPattern(ge_op.getIndex(), m_Constant(&index))) {
      return failure();
    }
    auto preceding_build_list = llvm::dyn_cast_or_null<BuildListOp>(
        ge_op.getTensorList().getDefiningOp());