// Add TF->TF lowering patterns.
  TF::PopulateTFLoweringBeforeHLOPatterns(context, &legalize_lower_patterns);

  if (tf2xla_fallback_device_type && prefer_tf2xla) {
    VLOG(1) << "TF to XLA legalization patterns are partitioned by op into "
               "either native MLIR legalization, or TF2XLA fallback "
               "legalzation, with a preference toward TF2XLA.";
  } else if (tf2xla_fallback_device_type) {

      FunctionType::get(func.getContext(), new_input_types,
                        func.front().getTerminator()->getOperandTypes()));
}

// Contains cached information for decomposed callee functions for (stateful)
// partitioned call ops.
struct PartitionedCallStackOpsInfo {
  bool signature_change;
  func::FuncOp decomposed_callee;
  llvm::SmallDenseMap<int64_t, int64_t> stack_var_arg_to_size_arg;
};

std::unique_ptr<OperationPass<ModuleOp>> CreateTPUResourceReadForWritePass();

// Creates a pass that reorders partitiioned resource reads and replicated
// inputs.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateTPUReorderReplicateAndPartitionedInputsPass();

// Creates a pass that partitions unpartitioned resource read/write to
// partitioned resource variables.
std::unique_ptr<OperationPass<func::FuncOp>>

// in this function should be on the same host.
struct FunctionMetadata {
  // The original function name before partition.
  llvm::StringRef original_name;
  // The insertion point of partition functions.
  Block::iterator insertion_point;
  // The partitioned function name.
  llvm::StringRef partition_name;
  // The input values of the function.
  llvm::SmallVector<Value, 4> inputs;

func.func @cluster_func(%arg0: tensor<*xi32>) -> tensor<*xi32> {
  func.return %arg0 : tensor<*xi32>
}

// -----

// Tests output sharding of unpartitioned resource write takes on same sharding
// as unpartitioned resource.

// CHECK-LABEL: func @partitioned_input_output
func.func @partitioned_input_output(%arg0: tensor<!tf_type.resource<tensor<f32>>>) {

// which returns a dict of string key -> EITHER:
// 1. python callable (taking a checkpoint key) returning SaveableObject OR
// 2. variable (partitioned/resource/reference or otherwise)
// https://github.com/tensorflow/tensorflow/blob/ddc1bbad3dfd4a089eb96014f26cc16664b1b2f8/tensorflow/python/training/saving/saveable_object.py#L58.

  }

  /**
   * Guides the allowed concurrency among update operations. Used as a hint for internal sizing. The
   * table is internally partitioned to try to permit the indicated number of concurrent updates
   * without contention. Because assignment of entries to these partitions is not necessarily
   * uniform, the actual concurrency observed may vary. Ideally, you should choose a value to

  }

  /**
   * Guides the allowed concurrency among update operations. Used as a hint for internal sizing. The
   * table is internally partitioned to try to permit the indicated number of concurrent updates
   * without contention. Because assignment of entries to these partitions is not necessarily
   * uniform, the actual concurrency observed may vary. Ideally, you should choose a value to

// RUN: tf-opt %s -split-input-file -verify-diagnostics -tf-tpu-reorder-replicate-partitioned-inputs | FileCheck %s

// CHECK-LABEL:func @simple
// CHECK-SAME: ([[ARG0:%.*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>, [[ARG1:%.*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>, [[ARG2:%.*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>, [[ARG3:%.*]]: tensor<!tf_type.resource<tensor<10x3xf32>>>)

        from = partitionPoint + 1;
      }
    }
  }

  /**
   * Performs a partition operation on the slice of {@code array} with elements in the range [{@code
   * from}, {@code to}]. Uses the median of {@code from}, {@code to}, and the midpoint between them
   * as a pivot. Returns the index which the slice is partitioned around, i.e. if it returns {@code