const xla::Shape& on_device_shape,
                                       xla::LocalClient* client,
                                       int device_ordinal) {
  xla::Shape on_host_shape =
      xla::ShapeUtil::DeviceShapeToHostShape(on_device_shape);
  xla::ScopedShapedBuffer shaped_buffer(on_host_shape, on_device_shape,
                                        client->backend().memory_allocator(),

  // shape. If a ShapedBuffer exists already (has_shaped_buffer() == true), it
  // is replaced and the managed memory deallocated.
  Status AllocateShapedBuffer(DataType dtype, const xla::Shape& on_device_shape,
                              xla::LocalClient* client, int device_ordinal);

  // Some Tensors can have complex on-device shapes, including tuple shapes. To