" (", num_tasks, " * ", num_tpus_per_task, " * ", kTPUTopologyRank,
        "), got ", topology_proto.device_coordinates_size()));

  const int bound_x = topology_proto.mesh_shape(0);
  const int bound_y = topology_proto.mesh_shape(1);
  const int bound_z = topology_proto.mesh_shape(2);
  const int bound_core = topology_proto.mesh_shape(3);

  xla::Array4D<TaskAndDevice> topology(bound_x, bound_y, bound_z, bound_core);
  int pos = 0;

string ArgView::VariableName() const { return arg_.name(); }

string ArgView::SetterMethod() const {
  if (IsList()) {
    return "AddInputList";
  } else {
    return "AddInput";
  }
}

std::vector<string> ArgView::SetterArgs() const { return {VariableName()}; }

bool ArgView::IsList() const { return arg_.arg_type().is_list(); }

int ArgView::Position() const { return arg_.position(); }

  // pass the filter (returning true) will be included.
  const llvm::SmallVector<Operation*, 4>& DirectControlPredecessors(
      Operation* op) const;
  llvm::SmallVector<Operation*, 4> DirectControlPredecessors(
      Operation* op, llvm::function_ref<bool(Operation*)> filter) const;

  // pass the filter (returning true) will be included.
  const llvm::SmallVector<Operation*, 4>& DirectControlSuccessors(

  // IsUnknownResource(resource) == false.
  const llvm::SmallSet<int64_t, 8>& GetResourceUniqueIds(Value resource) const;

  // Returns the set of values that are potentially aliases of `value`. Requires
  // `IsUnknownResource(resource) == false`.
  llvm::SmallSetVector<Value, 8> GetResourceAliases(Value resource) const;

  llvm::SmallSetVector<Value, 8> GetValuesForResourceId(int64_t id) const {
    auto it = id_to_resource_values_.find(id);

namespace tensorflow {
namespace generator {

string Call(const string& object, const string& method,
            std::vector<string> arguments, const char* oper) {
  return absl::Substitute("$0$1$2($3)", object, oper, method,
                          absl::StrJoin(arguments, ", "));
}

string Call(const string& function, std::vector<string> arguments) {

  return absl::OkStatus();
}

absl::Status ValidateSPAllocatorStats(const SP_AllocatorStats& stats) {
  TF_VALIDATE_STRUCT_SIZE(SP_AllocatorStats, stats,
                          SP_ALLOCATORSTATS_STRUCT_SIZE);
  // All other fields could theoretically be zero/null.
  return absl::OkStatus();
}

absl::Status ValidateSPDeviceMemoryBase(const SP_DeviceMemoryBase& mem) {
  TF_VALIDATE_STRUCT_SIZE(SP_DeviceMemoryBase, mem,

    std::string has_side_effect_key = "has_side_effect";
    bool has_side_effect_set = false;
    const flexbuffers::Map& computation_map =
        flexbuffers::GetRoot(op->custom_attributes).AsMap();
    const auto& keys = computation_map.Keys();
    for (size_t i = 0; i < keys.size(); ++i) {
      const auto key = keys[i].AsKey();
      const auto& value = computation_map[key];
      if (has_side_effect_key == key) has_side_effect_set = true;

// enable logging.
inline const char kBridgeComponent[] = "TFXLABridge";
inline const char kMlirPh1BridgeCounterReplicated[] = "replicated";
inline const char kMlirPh1BridgeCounterNonReplicated[] = "nonreplicated";
inline const char kMlirPh1BridgeCounterV1[] = "v1";
inline const char kMlirPh1BridgeCounterV2[] = "v2";
inline const char kMlirPh1BridgeCounterTpu[] = "tpu";
inline const char kMlirPh1BridgeCounterNonTpu[] = "cpu/gpu";

namespace generator {

string SourceCode::Render() const {
  string code;
  for (const Line& line : lines_) {
    absl::StrAppend(&code, string(line.indent * spaces_per_indent_, ' '),
                    line.text, "\n");
  }
  return code;
}

void SourceCode::AddLineWithIndent(const string& line) {
  ValidateAndAddLine(current_indent_, line);
}

void SourceCode::AddLineWithoutIndent(const string& line) {
  ValidateAndAddLine(0, line);

  // CHECK:      %[[ELEM_0:.*]] = "tf.Const"()
  // CHECK:      {ici_weight_distribution_mlir_bridge_marker = true}
  // CHECK-NEXT: %[[SHAPE_0:.*]] = "tf.Const"()
  // CHECK:      {ici_weight_distribution_mlir_bridge_marker = true}
  // CHECK-NEXT: %[[FULL_0:.*]] = "tf.Fill"(%[[SHAPE_0]], %[[ELEM_0]]) {ici_weight_distribution_mlir_bridge_marker = true}
  // CHECK:      %[[ELEM_1:.*]] = "tf.Const"()