op_name == "TruncatedNormal" || op_name == "Multinomial";
  }

  bool OpProducesOrConsumesVariant(const Node& node) const {
    auto is_variant = [](DataType dtype) { return dtype == DT_VARIANT; };
    return absl::c_any_of(node.input_types(), is_variant) ||
           absl::c_any_of(node.output_types(), is_variant);
  }

  bool HasXLAKernel(const Node& node,

  tensorflow::TensorShape shape(variable.shape());
  tensorflow::DataType dtype = variable.dtype();
  std::vector<std::string> component_devices;

  for (const auto& component :
       variable.experimental_distributed_variable_components()) {
    component_devices.push_back(component.device());
  }

  TF_RETURN_IF_ERROR(Variable::CreateUninitialized(
      ctx, dtype, shape, name,

  virtual AbstractTensorInterface* CreateBoolScalar(bool value) = 0;

  // Tensor creation functions
  virtual AbstractTensorInterface* CreateTensor(
      DataType dtype, absl::Span<const int64_t> dim_sizes) = 0;

  typedef void (*MemoryReleaser)(void* data, size_t len, void* arg);

  // Create a tensor instance from the given data buffer and description.

  // pushes and pops at the back.
  std::stack<AccumulatorCallState> call_state_;
};

// Template instantiations here

inline bool IsDtypeTrainable(DataType dtype) {
  switch (dtype) {
    case DT_HALF:
    case DT_BFLOAT16:
    case DT_FLOAT:
    case DT_DOUBLE:
    case DT_COMPLEX64:
    case DT_COMPLEX128:
    case DT_RESOURCE:
    case DT_VARIANT:

  }

  tensorflow::AbstractTensorInterface* t =
      tensorflow::unwrap(ctx)->CreateTensor(
          static_cast<tensorflow::DataType>(dtype), dimvec);

  if (t == nullptr) {
    status->status =
        tensorflow::errors::InvalidArgument("Unsupported dtype: ", dtype);
    return nullptr;
  }

  return new TF_Tensor{t};
}

TFE_TensorHandle* TFE_NewTensorHandleFromTensor(TFE_Context* ctx, TF_Tensor* t,

Status DefaultPaddedShapeFn(const Tensor& tensor, xla::Shape* shape) {
  const tensorflow::XlaTensor* xla_tensor =
      tensorflow::XlaTensor::FromTensor(&tensor);
  if (xla_tensor == nullptr) {
    return TensorShapeToXLAShape(tensor.dtype(), tensor.shape(), shape);
  }

  const xla::ShapedBuffer& shaped_buffer = xla_tensor->shaped_buffer();
  *shape = shaped_buffer.on_device_shape();
  return absl::OkStatus();
}

    "'h_scale' type: 'float'} attr : { name: 'max_classes_per_detection' "
    "type: 'int'} attr : { name: 'max_detections' type: 'int'} attr : { "
    "name: 'nms_iou_threshold' type: 'float'} attr : { name: "
    "'nms_score_threshold' type: 'float'} attr : { name: 'num_classes' type: "
    "'int'} attr : { name: 'w_scale' type: 'float'} attr : { name: 'x_scale' "

  CreateContext();

  std::vector<XlaCompiler::Argument> args(2);
  args[0].kind = XlaCompiler::Argument::kParameter;
  args[0].type = DT_INT32;
  args[0].shape = TensorShape({1, 3});
  args[1].kind = XlaCompiler::Argument::kParameter;
  args[1].type = DT_INT32;
  args[1].shape = TensorShape({1, 3});

  const XlaCompiler::CompilationResult* result;
  xla::PjRtLoadedExecutable* executable;

    TF_Status* status);

// Information about the shape of a Tensor and its type.
struct TF_ShapeAndType {
  // Number of dimensions. -1 indicates unknown rank.
  int num_dims;
  // Array of dimensions. -1 indicates unknown dim.
  int64_t* dims;
  // The data type. May be 0 to denote unknown type.
  TF_DataType dtype;
};

typedef struct TF_ShapeAndType TF_ShapeAndType;

         << compilation_result.xla_output_shape.ToString() << '\n';

  for (const auto& xla_output_description : compilation_result.outputs) {
    output << "// XlaOutputDescription type="
           << DataTypeString(xla_output_description.type) << " shape=("
           << absl::StrJoin(xla_output_description.shape.dim_sizes(), ", ")
           << ')';
    if (xla_output_description.input_index >= 0)