APIs
PyTorch Native APIs¶
Most of PyTorch native functions can be used on RBLN NPU.
RBLN Specific APIs¶
The following functions are RBLN NPU-specific functions defined in the torch.rbln module.
Classes¶
device_of
¶
Context-manager that changes the current device to that of given object.
You can use both tensors and storages as arguments. If a given object is not allocated on an RBLN device, this is a no-op.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
obj
|
Tensor or Storage
|
object allocated on the selected device. |
required |
Functions¶
get_amp_supported_dtype()
¶
Get a list of data types supported by automatic mixed precision (AMP) on RBLN devices.
Returns:
| Type | Description |
|---|---|
List[dtype]
|
List[torch.dtype]: A list of data types supported by AMP. |
Note
This function currently returns only torch.float16.
is_available()
¶
Check if any RBLN devices are available.
Returns:
| Name | Type | Description |
|---|---|---|
bool |
bool
|
True if at least one RBLN device is available, False otherwise. |
current_device()
¶
Get the index of the currently selected RBLN device.
Returns:
| Name | Type | Description |
|---|---|---|
int |
int
|
The index of the currently selected RBLN device. |
device_count()
¶
Get the number of available RBLN devices.
Returns:
| Name | Type | Description |
|---|---|---|
int |
int
|
The number of available RBLN devices. |
set_device(device)
¶
Set the current device.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
device
|
device or int or str
|
selected device. |
required |