This document is relevant for: Inf2, Trn1, Trn2, Trn3

nrtpy tensor API#

This page documents NrtpyTensor — the device-resident tensor with NumPy integration and automatic cleanup.

NrtpyTensor#

class nrtpy.NrtpyTensor(tensor_ref, shape, dtype, name=None)#

A tensor resident in high bandwidth memory (HBM) on the device. Device memory is freed automatically when the tensor is garbage collected by Python. For explicit control, call nrtpy.get_nrtpy_singleton().free_tensor(tensor)().

Parameters:

  • tensor_ref (nrtpy.NrtTensor) – Reference to the underlying device tensor.

  • shape (tuple[int, ...] or int) – Shape of the tensor. An integer is treated as a single-dimension shape.

  • dtype (numpy.dtype) – Data type of the tensor, in NumPy format.

  • name (str or None) – Optional name for the tensor.

Attributes

shape: tuple[int, ...]#

Shape of the tensor.

dtype: numpy.dtype#

Data type of the tensor.

name: str#

Name of the tensor.

classmethod from_numpy(array, name=None, core_id=0)#

Allocate a device tensor and copy a NumPy array into it.

Parameters:

  • array (numpy.ndarray) – Source array. It is made contiguous before the copy.

  • name (str or None) – Optional name for the tensor.

  • core_id (int) – Target NeuronCore for the allocation (default 0).

Returns:

An NrtpyTensor backed by the newly allocated device memory.

Return type:

nrtpy.NrtpyTensor

write_from_numpy(array)#

Write new data from a NumPy array into this existing device tensor without reallocating.

Parameters:

array (numpy.ndarray) – Source array. It is made contiguous before the copy and must match the tensor’s byte size.

Raises:

ValueError – If the source array’s byte size does not match the tensor’s byte size.

numpy()#

Read the tensor data back from the device as a NumPy array.

Returns:

A NumPy array with the tensor’s shape and dtype.

Return type:

numpy.ndarray

Note

The float8_e4m3 and float8_e5m2 dtypes are reported as int8 by libnrt. nrtpy includes workarounds, but be aware of this when inspecting dtypes on FP8 tensors.

Examples#

Create, write, and read#

import numpy as np
from nrtpy import NrtpyTensor

# Allocate a device tensor from a NumPy array
data = np.random.randn(1, 10).astype(np.float32)
tensor = NrtpyTensor.from_numpy(data, name="input")

# Overwrite in place (no reallocation)
new_data = np.zeros((1, 10), dtype=np.float32)
tensor.write_from_numpy(new_data)

# Read back to host
host_array = tensor.numpy()

Multiple cores#

import nrtpy
from nrtpy import NrtpyTensor

nrtpy.configure(visible_cores=[0, 1])

# Allocate tensors on different cores
t0 = NrtpyTensor.from_numpy(data, name="core0_input", core_id=0)
t1 = NrtpyTensor.from_numpy(data, name="core1_input", core_id=1)