zarr.buffer.gpu#

Attributes#

buffer_prototype

Classes#

`Buffer`	A flat contiguous memory block on the GPU
`NDBuffer`	A n-dimensional memory block on the GPU

Module Contents#

class zarr.buffer.gpu.Buffer(array_like: zarr.core.buffer.core.ArrayLike)[source]#

Bases: zarr.core.buffer.core.Buffer

A flat contiguous memory block on the GPU

We use Buffer throughout Zarr to represent a contiguous block of memory.

A Buffer is backed by a underlying array-like instance that represents the memory. The memory type is unspecified; can be regular host memory, CUDA device memory, or something else. The only requirement is that the array-like instance can be copied/converted to a regular Numpy array (host memory).

Parameters:

array_like: array-like object that must be 1-dim, contiguous, and byte dtype.

Notes

This buffer is untyped, so all indexing and sizes are in bytes.

as_array_like() → ArrayLike[source]#

Returns the underlying array (host or device memory) of this buffer

This will never copy data.

Returns:

The underlying 1d array such as a NumPy or CuPy array.

as_buffer_like() → zarr.core.common.BytesLike[source]#

Returns the buffer as an object that implements the Python buffer protocol.

Returns:

An object that implements the Python buffer protocol

Notes

Might have to copy data, since the implementation uses .as_numpy_array().

as_numpy_array() → numpy.typing.NDArray[Any][source]#

Returns the buffer as a NumPy array (host memory).

Returns:

NumPy array of this buffer (might be a data copy)

Notes

Might have to copy data, consider using .as_array_like() instead.

classmethod create_zero_length() → Self[source]#

Create an empty buffer with length zero

Returns:

New empty 0-length buffer

classmethod from_array_like(array_like: ArrayLike) → Self[source]#

Create a new buffer of an array-like object

Parameters:

array_like: array-like object that must be 1-dim, contiguous, and byte dtype.

Returns:

New buffer representing array_like

classmethod from_buffer(buffer: zarr.core.buffer.core.Buffer) → Self[source]#

Create an GPU Buffer given an arbitrary Buffer This will try to be zero-copy if buffer is already on the GPU and will trigger a copy if not.

Returns:

New GPU Buffer constructed from buffer

classmethod from_bytes(bytes_like: zarr.core.common.BytesLike) → Self[source]#

Create a new buffer of a bytes-like object (host memory)

Parameters:

bytes_like: bytes-like object

Returns:

New buffer representing bytes_like

to_bytes() → bytes[source]#

Returns the buffer as bytes (host memory).

Returns:

bytes of this buffer (data copy)

Warning

Will always copy data, only use this method for small buffers such as metadata buffers. If possible, use .as_numpy_array() or .as_array_like() instead.

class zarr.buffer.gpu.NDBuffer(array: zarr.core.buffer.core.NDArrayLike)[source]#

Bases: zarr.core.buffer.core.NDBuffer

A n-dimensional memory block on the GPU

We use NDBuffer throughout Zarr to represent a n-dimensional memory block.

A NDBuffer is backed by a underlying ndarray-like instance that represents the memory. The memory type is unspecified; can be regular host memory, CUDA device memory, or something else. The only requirement is that the ndarray-like instance can be copied/converted to a regular Numpy array (host memory).

Parameters:

array: ndarray-like object that is convertible to a regular Numpy array.

Notes

The two buffer classes Buffer and NDBuffer are very similar. In fact, Buffer is a special case of NDBuffer where dim=1, stride=1, and dtype=”B”. However, in order to use Python’s type system to differentiate between the contiguous Buffer and the n-dim (non-contiguous) NDBuffer, we keep the definition of the two classes separate.

all_equal(other: Any, equal_nan: bool = True) → bool[source]#: Compare to other using np.array_equal.

as_ndarray_like() → NDArrayLike[source]#

Returns the underlying array (host or device memory) of this buffer

This will never copy data.

Returns:

The underlying array such as a NumPy or CuPy array.

as_numpy_array() → numpy.typing.NDArray[Any][source]#

Returns the buffer as a NumPy array (host memory).

Returns:

NumPy array of this buffer (might be a data copy)

Warning

Might have to copy data, consider using .as_ndarray_like() instead.

as_scalar() → ScalarType[source]#: Returns the buffer as a scalar value

astype( dtype: numpy.typing.DTypeLike, order: Literal['K', 'A', 'C', 'F'] = 'K', ) → Self[source]#

copy() → Self[source]#

classmethod create( *, shape: collections.abc.Iterable[int], dtype: numpy.typing.DTypeLike, order: Literal['C', 'F'] = 'C', fill_value: Any | None = None, ) → Self[source]#

Create a new buffer and its underlying ndarray-like object

Parameters:

shape: The shape of the buffer and its underlying ndarray-like object
dtype: The datatype of the buffer and its underlying ndarray-like object
order: Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory.
fill_value: If not None, fill the new buffer with a scalar value.

Returns:

New buffer representing a new ndarray_like object

Notes

A subclass can overwrite this method to create a ndarray-like object other then the default Numpy array.

classmethod empty( shape: tuple[int, Ellipsis], dtype: numpy.typing.DTypeLike, order: Literal['C', 'F'] = 'C', ) → Self[source]#

Create an empty buffer with the given shape, dtype, and order.

This method can be faster than NDBuffer.create because it doesn’t have to initialize the memory used by the underlying ndarray-like object.

Parameters:

shape: The shape of the buffer and its underlying ndarray-like object
dtype: The datatype of the buffer and its underlying ndarray-like object
order: Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory.

Returns:

buffer: New buffer representing a new ndarray_like object with empty data.