} else {
    return false;
  }
  return true;
}

// Create a dummy zero to be fed locally from the host to the TPUExecute.
Value CreateZeroInput(Location loc, OpBuilder& builder, Attribute zero_attr,
                      DenseIntElementsAttr shape_attr) {
  ConstOp zero = builder.create<ConstOp>(loc, zero_attr);
  zero->setAttr(kICIWeightDistributionMlirBridgeMarker,

  builder.setInsertionPoint(&terminator);
  builder.create<tf_executor::YieldOp>(terminator.getLoc(),
                                       terminator.getOperands());
  terminator.erase();

  builder.setInsertionPoint(island_op);
  IRMapping mapping;
  for (int i : llvm::seq<int>(0, num_replicas)) {
    // Create new island for replica.
    auto replica = builder.create<tf_executor::IslandOp>(

  if (subtypes.size() != 1) {
    return std::nullopt;
  }
  return subtypes[0].getElementType();
}

// Create an `ConstBytesAttr` which encodes the options
// for the `tf.custom` tensor list op to be created. If the given
// op is not a `tf.TensorList*` op, return empty, although this case
// should never be trigged in practice since patterns are only applied
// on `tf.TensorList*` ops.

  }

  auto dataset_while = builder.create<TF::WhileRegionOp>(
      reduce_dataset.getLoc(), while_input_types, /*input=*/while_input_values,
      /*parallel_iterations=*/10, false,
      /*shape_invariant=*/false);

  // `_lower_using_switch_merge` is the default for While ops created
  // in TensorFlow and allows lowering to V1 control flow for loop
  // parallelization.

      tensorflow::GetTypeFromTFTensorShape({1}, shape_dtype);
  auto expanded_start_index = rewriter->create<TF::ExpandDimsOp>(
      loc, vector_type, start_index, scalar_zero);
  auto start_position = rewriter->create<TF::ConcatOp>(
      loc, position_type, scalar_zero,
      ArrayRef<Value>({expanded_start_index, partial_start_position}));

  // Create the slice size tensor. This is done by concatenating `size` and
  // `partial_size`.

        xla_tensor->AllocateShapedBuffer(device_tensor->dtype(), shape, client_,
                                         stream_->parent()->device_ordinal()));

    // The cpu_tensor and literal that we created here hold the data of host
    // tensor in descending layout. The layout could be different from layout in
    // device_tensor (but the logical shape has to be the same). The

    TF_DeleteStatus(status);
  }

  // 6. Read `var` using `ctx_1`. This read should succeed since `ctx_1` was
  // created with `isolate_session_state` set to false, and update should
  // preserve it.
  {
    // Create a handle to `var`, using `ctx_1`.
    TFE_TensorHandle* var_handle =
        CreateVarHandle(ctx_1, remote_device, /*variable_name=*/"var");

          attr_cst =
              builder.create<ConstOp>(op->getLoc(), output_type, attribute);
        } else {
          attr_cst = builder.create<mlir::arith::ConstantOp>(
              op->getLoc(), cast<TypedAttr>(attribute));
        }
        new_operands.push_back(attr_cst);
      }
    }

    // Create the TFR call op
    auto new_op = builder.create<CallOp>(

    std::string tmp_dir{};
    if (!env_->LocalTempFilename(&tmp_dir)) {
      return absl::InternalError("Failed to create temp file.");
    }

    TF_CHECK_OK(env_->CreateDir(tmp_dir));
    return tmp_dir;
  }

  // Parses `module_op_str` to create a `ModuleOp`. Checks whether the created
  // module op is valid.
  mlir::OwningOpRef<mlir::ModuleOp> ParseModuleOpString(
      const absl::string_view module_op_str) {

      absl::StrCat("device: ", DEVICE_TPU_SYSTEM, " device"));
  devices->push_back(std::make_unique<VirtualDevice>(options.env, attrs));
  VLOG(1) << "Created TPU_SYSTEM device. This host has " << device_count
          << " TPUs";

  return absl::OkStatus();
}

}  // namespace

void RegisterTpuDeviceToDeviceCopy() {
  static auto* const register_tpu_tpu_copy = new CopyTensor::Registration(