// cluster to before the cluster.
mlir::LogicalResult LiftHeadOutsideCompiledOps(
    OpBuilder* builder,
    const mlir::TF::SideEffectAnalysis& side_effect_analysis,
    const mlir::TF::RuntimeDevices& devices, mlir::tf_device::ClusterOp cluster,
    std::string* host_device, bool* cluster_updated) {
  llvm::SmallVector<Operation*, 4> head_outside_compiled_ops =

      "SmallVector<unsigned, 4>", "GetLayoutDependentResults", (ins)
    >,
    InterfaceMethod<
      [{Returns the optimal data layout based on the available devices.}],
      "StringRef", "GetOptimalLayout", (ins "const RuntimeDevices&":$devices)
    >,
    InterfaceMethod<
      [{Updates operation attributes and operands to account for the updated
        data format. If data format is not supported, must return failure.}],

  if (!tf_dialect) {
    getOperation().emitError() << "'tf' dialect is not registered";
    return signalPassFailure();
  }

  // Get the number of devices per host.
  int device_num = 0;
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(
          getOperation()->getParentOfType<ModuleOp>(), &devices)))
    return signalPassFailure();
  for (const auto& device_name : devices.device_names()) {

    bool pass_host_device, ModuleOp module,
    std::function<WalkResult(Operation*, tf_device::ClusterOp,
                             std::optional<std::string>)>
        callback) {
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices))) return failure();
  const CallGraph call_graph(module);
  // symbol_table caches callees in the CallGraph.
  SymbolTableCollection symbol_table;

}

LogicalResult GetTpuDeviceAssignment(
    ClusterOp cluster, ReplicateOp replicate, mlir::ModuleOp module,
    absl::StatusOr<TPUDeviceAssignment>& status_or_tpu_device_assignment) {
  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices))) return failure();

  uint32_t num_replicas = replicate.getN();

  auto num_cores_per_replica_attr = cluster->getAttrOfType<mlir::IntegerAttr>(

  }
  return success();
}

LogicalResult BiasAddOp::UpdateDataFormat(StringRef data_format) {
  return ::mlir::TF::UpdateDataFormat(data_format, this);
}

StringRef BiasAddOp::GetOptimalLayout(const RuntimeDevices& devices) {
  // Keep current data format if no GPUs are available or if explicit placement
  // does not allow to use GPU for this operation.
  if (!CanUseGpuDevice(devices) || !CanUseGpuDevice(getOperation()))

using Permutation = SmallVector<int64_t, 4>;

void LayoutAssignmentPass::runOnOperation() {
  func::FuncOp func = getOperation();

  // Get runtime devices information from the closest parent module.
  RuntimeDevices devices;
  if (failed(::tensorflow::GetDevicesFromOp(func->getParentOfType<ModuleOp>(),
                                            &devices)))
    return signalPassFailure();

      new_parallel_execute, builder);
  if (failed(result)) return failure();

  return TF::RemoveSingletonParallelExecuteOp(new_parallel_execute, builder);
}

void TPURewritePass::runOnOperation() {
  TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(getOperation(), &devices)))
    return signalPassFailure();

  // Collect compilation results.

  // Get runtime devices information from the closest parent module.
  auto module = getOperation();
  if (failed(CheckPreconditions(module))) signalPassFailure();

  mlir::TF::RuntimeDevices devices;
  if (failed(tensorflow::GetDevicesFromOp(module, &devices)))
    return signalPassFailure();

  llvm::SmallVector<mlir::tf_device::ParallelExecuteOp, 4>
      tmp_parallel_execute_ops;

    // TF_LayoutSensitiveInterface:
    SmallVector<unsigned, 4> GetLayoutDependentArgs() { return {0}; }
    SmallVector<unsigned, 4> GetLayoutDependentResults() { return {0}; }
    StringRef GetOptimalLayout(const RuntimeDevices& devices);
    LogicalResult UpdateDataFormat(StringRef data_format);
  }];

  let hasVerifier = 1;
}

def TF_BiasAddGradOp : TF_Op<"BiasAddGrad", [Pure]> {
  let summary = [{