// Pack 3 variable handles into one TFE_TensorHandle.
  // When remote is false, function device is placed on task0. Handle types are
  // REMOTE, REMOTE, LOCAL on task0. When remote is true, function device is
  // placed on task1, Handle types are LOCAL, REMOTE, LOCAL on task1.
  int num_replicas = 3;
  std::vector<TFE_TensorHandle*> handles = {h0, h1, h2};
  TFE_TensorHandle* packed_handle =

                 absl::Span<AbstractTensorHandle* const> grad_outputs,
                 absl::Span<AbstractTensorHandle*> grad_inputs) override {
    /* Given upstream grad U and a matmul op A*B, the gradients are:
     *
     *    dA = U * B.T
     *    dB = A.T * U
     *
     *    where A.T means `transpose(A)`
     */
    AbstractTensorHandle* upstream_grad = grad_outputs[0];

    // Get transpose attrs

//
// This has three benefits:
//   * allows having newer plugins than the current core TensorFlow: the
//     additional entries in the plugin's table are just discarded;
//   * allows having older plugins than the current core TensorFlow (though
//     we are still warning users): the entries that core TensorFlow expects
//     but plugins didn't provide will be set to `nullptr` values and core

    Run(inputs, {{output, 0}}, {expected_result});
  }

  // Run the host graph, which now contains a function and check that
  // outputs are as expected.
  // 'T' stands for 'tensor' since the outputs are tensors, not scalars.
  void RunT(const std::vector<std::pair<TF_Operation*, TF_Tensor*>>& inputs,
            std::initializer_list<TF_Output> outputs,

    const char* fn_name = "test_fn";
    core::RefCountPtr<AbstractFunction> scoped_func;
    // Returning null tensors from a tf.function is not supported, so we keep
    // track of indices in the model's outputs are nullptr in this set.
    // The FunctionDef only outputs the non-null tensors. We later pad the
    // function op outputs to have nullptrs at the `null_indices`.
    absl::flat_hash_set<int> null_indices;
    {

    if (gcs_file->stat_cache->Lookup(path, &stat)) {
      if (offset + read < stat.base.length) {
        TF_SetStatus(status, TF_INTERNAL,
                     absl::StrCat("File contents are inconsistent for file: ",
                                  path, " @ ", offset)
                         .c_str());
      }
      TF_VLog(2, "Successful integrity check for: %s @ %u", path.c_str(),
              offset);

                    /*input_shapes*/ {make_shape({3, 2}), unknown_shape()},
                    /*input_tensors*/ {},
                    /*expected_shape*/ make_shape({3, kUnknownDim}));

  // Infer shape when some dimensions are unknown.
  CheckOutputShapes(
      matmul_op,
      /*input_shapes*/ {make_shape({kUnknownDim, 2}), make_shape({2, 4})},
      /*input_tensors*/ {},
      /*expected_shape*/ make_shape({kUnknownDim, 4}));

  }
  ~GraphOperation() override {}

 private:
  friend class GraphContext;  // For access to op_.
  TF_Graph* g_;
  std::unique_ptr<TF_OperationDescription> op_;
  // Hold `op_type` and `op_name` till both are available since we need both
  // to build a graph operation.
  string op_type_;
  const char* op_name_ = nullptr;
  // TODO(srbs): Use this.
  string device_name_;
};

// TODO(b/291142876) Simplify TFE_ContextSetServerDefWithTimeoutAndRetries and
// TFE_ContextUpdateServerDefWithTimeout to be simple wrappers around the same
// C++ function.
// Retries are used for CreateContext calls, which is used in
// ParameterServerStrategy initialization to be robust to worker preemption.
TF_CAPI_EXPORT extern void TFE_ContextSetServerDefWithTimeoutAndRetries(

    string key(v2_reader_->key());
    (*var_to_shape_map)[key] = TensorShape(entry.shape());
    (*var_to_data_type_map)[key] = DataType(entry.dtype());
  }
  // The returned pointers are owned by the caller.
  return std::make_pair(std::move(var_to_shape_map),
                        std::move(var_to_data_type_map));
}

}  // namespace checkpoint