TF_RETURN_IF_ERROR(
      fld->ReplaceFunction(shape_inference_graph_name, fdef_replace));

  return absl::OkStatus();
}

void SetMaximalSharding(NodeDefBuilder& node_builder) {
  xla::OpSharding sharding;
  sharding.set_type(xla::OpSharding::MAXIMAL);
  sharding.add_tile_assignment_dimensions(1);
  sharding.add_tile_assignment_devices(0);