absl::Span<const int64_t> shape, const TF_DataType dtype)
      : device_(device),
        tensors_(std::move(tensors)),
        shape_(std::vector<int64_t>(shape.begin(), shape.end())),
        dtype_(dtype) {}
  ParallelTensor(const ParallelDevice& device,
                 std::vector<TensorHandlePtr> tensors, const TF_DataType dtype)
      : device_(device),
        tensors_(std::move(tensors)),

  }
  TFE_TensorHandleGetStatus(components[0].get(), status);
  if (!status->status.ok()) {
    return nullptr;
  }

  TF_DataType dtype = TFE_TensorHandleDataType(components[0].get());
  // Verify that the TensorHandle's shape and dtype match all of the component
  // shapes and dtypes.
  for (TensorHandlePtr& component : components) {
    TFE_TensorHandleGetStatus(component.get(), status);
    if (!status->status.ok()) {

#define TENSORFLOW_C_EAGER_IMMEDIATE_EXECUTION_CONTEXT_H_

#include <functional>
#include <map>
#include <memory>
#include <vector>

#include "absl/types/optional.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/immediate_execution_distributed_manager.h"
#include "tensorflow/c/eager/immediate_execution_operation.h"

  owner->reference.Unref();
  delete owner;
}

// Converts TF_DATAType to DLPack data type.
DLDataType GetDlDataType(TF_DataType data_type, TF_Status* status) {
  DLDataType dtype;
  dtype.lanes = 1;
  dtype.bits = TF_DataTypeSize(data_type) * 8;
  switch (data_type) {
    case TF_DataType::TF_BOOL:
      dtype.code = DLDataTypeCode::kDLBool;
      break;
    case TF_DataType::TF_HALF:
    case TF_DataType::TF_FLOAT:

  absl::Status status;
  if (reader_ != nullptr) {
    status = reader_->GetTensor(name, out_tensor);
  } else {
    tensorflow::DataType dtype;
    tensorflow::TensorShape shape;
    status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape);
    if (status.ok()) {
      out_tensor->reset(new Tensor(dtype, shape));
      status = v2_reader_->Lookup(name, out_tensor->get());
      if (!status.ok()) out_tensor->reset();
    }
  }

    if not fields:
        return fields
    first_field = fields[0]
    if len(fields) == 1 and lenient_issubclass(first_field.type_, BaseModel):
        fields_to_extract = get_cached_model_fields(first_field.type_)
        return fields_to_extract
    return fields


def get_flat_params(dependant: Dependant) -> List[ModelField]:

                               tensorflow::CustomDevice* device,
                               tensorflow::DataType dtype, void* data,
                               TFE_CustomDeviceTensorHandleMethods methods)
      : tensorflow::CustomDeviceTensorHandle(context, device, dtype),
        data_(data),
        methods_(methods) {}

  ~CAPICustomDeviceTensorHandle() override { methods_.deallocator(data_); }

  // Returns tensor (full) type.
  // While there is no immediate plan to deprecate dtype and shape in favor
  // of only using full type type information, this is a future possibility.
  //
  // Note that map_dtype_to_child_of_tensor() from core/framework/types.h
  // can be used to set a FullTypeDef based on dtype in a derived class if
  // appropriate.

        field_info = field.field_info
        if not (
            field.shape == SHAPE_SINGLETON  # type: ignore[attr-defined]
            and not lenient_issubclass(field.type_, BaseModel)
            and not lenient_issubclass(field.type_, dict)
            and not field_annotation_is_sequence(field.type_)
            and not is_dataclass(field.type_)
            and not isinstance(field_info, params.Body)
        ):
            return False

                                                void* data, size_t len,
                                                TF_Status* status) {
  auto dtype = static_cast<tensorflow::DataType>(data_type);
  DCHECK(tensorflow::DataTypeCanUseMemcpy(dtype));

  tensorflow::Tensor tensor(dtype, tensorflow::TensorShape({}));
  std::memcpy(tensorflow::TensorCApi::Buffer(tensor)->data(), data, len);

  status->status = absl::OkStatus();