// should be equal among all checkpoints written per job.
auto* checkpoint_size = monitoring::Counter<2>::New(
    "/tensorflow/core/checkpoint/write/checkpoint_size",
    "Size of checkpoint (.index and sharded data files), rounded to the "
    "nearest 100 MB.",
    "api_label", "filesize");

}  // namespace

// Counter that records how long it took to execute the checkpoint sharding
// callback in microseconds.

      return mlir::failure();
    shard_output_types.push_back(shard_type);
    full_output_types.push_back(output.getType());
  }

  // Convert split sharded inputs to MANUAL sharded inputs.
  // common_split_sharding is the split sharding that is common to all inputs
  // and outputs.
  llvm::SmallVector<Value, 4> manual_inputs;
  manual_inputs.reserve(inputs.size());

#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/platform/types.h"

namespace tensorflow {

namespace {
// All the passes we will make available to Python by default.
// TODO(tf): this should be sharded instead of being monolithic like that.
static void RegisterPasses() {
  static bool unique_registration = [] {
    mlir::registerAllPasses();
    mlir::registerTensorFlowPasses();

}

def TPUShardingIdentificationPass : Pass<"tf-tpu-sharding-identification", "ModuleOp"> {
  let summary = "Identifies and handles inputs/outputs of TPU computation that is "
           "sharded across logical cores.";
  let constructor = "tensorflow::tf2xla::internal::CreateTPUShardingIdentificationPass()";
  let description = [{
    Bubbles up sharding configuration from `cluster_func` regions into

by type `p`. There are exactly `GOMAXPROCS` Ps. A P can be thought of
like a CPU in the OS scheduler and the contents of the `p` type like
per-CPU state. This is a good place to put state that needs to be
sharded for efficiency, but doesn't need to be per-thread or
per-goroutine.

The scheduler's job is to match up a G (the code to execute), an M
(where to execute it), and a P (the rights and resources to execute

// runtime. This set of stats is grouped together because they
// depend on each other in some way to make sense of the runtime's
// current heap memory use. They're also sharded across Ps, so it
// makes sense to grab them all at once.
type heapStatsAggregate struct {
	heapStatsDelta

	// Derived from values in heapStatsDelta.

	// inObjects is the bytes of memory occupied by objects,

      producer: 23
    }
  }
  saver_def {
    filename_tensor_name: "save/Const:0"
    save_tensor_name: "save/Identity:0"
    restore_op_name: "save/restore_all"
    max_to_keep: 5
    sharded: true
    keep_checkpoint_every_n_hours: 10000.0
    version: V2
  }
  collection_def {
    key: "asset_filepaths"
    value {
      node_list {
        value: "Const:0"
      }
    }

  for (mlir::TF::IdentityOp to_erase : erase_list) {
    to_erase->erase();
  }

  return result_op;
}

// Return the cluster's per-replica result type, converting any full-shaped
// tensor types into sharded-shaped ones if they're partitioned.
llvm::SmallVector<Type, 8> GetClusterResultTypes(
    mlir::tf_device::ClusterOp cluster,
    const PartitionedClusterOutputMap& partitioned_outputs) {

  let hasVerifier = 1;
}

def TF_TPUPartitionedOutputOp : TF_Op<"TPUPartitionedOutput", [Pure]> {
  let summary = [{
An op that demultiplexes a tensor to be sharded by XLA to a list of partitioned
  }];

  let description = [{
outputs outside the XLA computation.
  }];

  let arguments = (ins
    TF_Tensor:$inputs,

    %1 = "tf.I"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    func.return %1 : tensor<?xi32>
  }
}

// -----

// Test `tf_device.cluster_func` on TPU with pre-split replicate sharded
// input/output using `tf.TPUPartitionedInputV2` and `tf.TPUPartitionedOutputV2`.