abracudabra.device.query

abracudabra.device.query#

Query the device of a NumPy/CuPy array, series or Torch tensor.

Functions#

frame_get_device_type(…)

Get the device type of a pandas or cudf series or dataframe.

get_device(…)

Get the device of a NumPy/CuPy array or series.

guess_device(*elements[, skip_unknown])

Guess the device of a sequence of arrays or tensors.

Module Contents#

abracudabra.device.query.frame_get_device_type(frame: abracudabra._annotations.Series | abracudabra._annotations.DataFrame, /, *, raise_if_unknown: Literal[True] = ...) → abracudabra.device.base.DeviceType[source]#
abracudabra.device.query.frame_get_device_type(frame: abracudabra._annotations.Series | abracudabra._annotations.DataFrame, /, *, raise_if_unknown: bool = ...) → abracudabra.device.base.DeviceType | None

Get the device type of a pandas or cudf series or dataframe.

Parameters:

  • frame – The frame to check.

  • raise_if_unknown – Whether to raise an error if the frame is not a series or dataframe.

Returns:

The device type of the frame. If raise_if_unknown is False and the frame is not a series or dataframe, return None.

Raises:

TypeError – If the frame is not a series or dataframe and raise_if_unknown is True.

abracudabra.device.query.get_device(element: abracudabra._annotations.Array | torch.Tensor, /, *, raise_if_unknown: Literal[True] = ...) → abracudabra.device.base.Device[source]#
abracudabra.device.query.get_device(element: abracudabra._annotations.Array | torch.Tensor, /, *, raise_if_unknown: bool = ...) → abracudabra.device.base.Device | None

Get the device of a NumPy/CuPy array or series.

Parameters:

  • element – The element to check.

  • raise_if_unknown – Whether to raise an error if the element is not a known array or tensor.

Returns:

The device of the element.

Examples

>>> import numpy as np
>>> array = np.random.rand(3)
>>> get_device(array)
Device(type="cpu", idx=None)
>>> import torch
>>> tensor = torch.rand(3, device="cuda")
>>> get_device(tensor)
Device(type="cuda", idx=0)
abracudabra.device.query.guess_device(*elements, skip_unknown=True)[source]#

Guess the device of a sequence of arrays or tensors.

This function checks the device of the elements and returns the device if all elements are on the same device. Otherwise, it raises an error.

Parameters:

  • *elements (abracudabra._annotations.Array | torch.Tensor) – The elements to check.

  • skip_unknown (bool) – Whether to skip elements that are not known arrays or tensors (e.g., lists, tuples, etc.).

Returns:

The device of the elements.

Raises:

  • ValueError – If no elements are given.

  • ValueError – If the device cannot be inferred from any of the elements (if skip_unknown is True).

  • ValueError – If the elements are on different devices.

Return type:

abracudabra.device.base.Device

Contents