nki.language.NkiTensor

This document is relevant for: Trn2, Trn3

nki.language.NkiTensor#

class nki.language.NkiTensor(shape: tuple[int, ...], dtype: str, storage: Any, buffer: nki.language.MemoryRegion | str = Ellipsis, name: str = '', address=None)[source]#

NKI tensor with shape-based view operations.

NkiTensor is the core tensor type in NKI. It represents a multi-dimensional array allocated on a specific memory buffer (SBUF, PSUM, or HBM) with a dtype and shape. All view operations (slice, permute, reshape, etc.) return new NkiTensor objects that share the same underlying storage — no data is copied. Views are consumed by NKI ISA instructions (e.g., nisa.tensor_copy, nisa.dma_copy).

Tensors are created via nl.ndarray():

sb = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
hbm = nl.ndarray((4, 128, 64), dtype=nl.float32, buffer=nl.shared_hbm)

Partition dimension.

On-chip tensors (SBUF, PSUM) have a partition dimension at dim 0 that maps to the hardware’s parallel partitions. Most view operations cannot modify this dimension — see individual method docs for constraints.

Methods

`ap`	Low-level access pattern override (escape hatch).
`broadcast`	Expand a size-1 dimension to `size` by repeating elements.
`expand_dim`	Insert a new dimension of size 1 at position `dim`.
`flatten_dims`	Merge a contiguous range of dimensions into one.
`get_pattern`	Return the view's access pattern as `[[stride, count], ...]`.
`indirect`	Create an indirect tensor view for Tensor Indirection (TI).
`is_contiguous`	Return True if the view covers storage contiguously (row-major order).
`is_indirect`	Return True if this view already uses indirect addressing.
`permute`	Reorder tensor dimensions.
`rearrange`	Rearrange tensor dimensions using einops-style patterns.
`reinterpret_cast`	Alias for `view()`.
`reshape`	Reshape the tensor to a new shape without copying data.
`reshape_dim`	Split a single dimension into multiple dimensions.
`select`	Select a single element along a dimension, removing it.
`slice`	Slice along a single dimension.
`squeeze_dim`	Remove a dimension of size 1.
`vector_select`	Per-partition indirect addressing using a vector of offsets.
`view`	Reinterpret the tensor's data as a different dtype.

Attributes

`buffer`	The memory region the storage lives in.
`dtype`	Element type, e.g.
`name`	Optional debug label, used in compiler diagnostics and the profiler.
`ndim`	Number of dimensions (`len(shape)`).
`offset`	Element offset (not byte offset) from the start of the underlying storage at which this view begins.
`shape`	Tuple giving the count of each dimension of the view.
`size`	Total number of elements (`prod(shape)`).
`strides`	Tuple with stride per dimension.

Low-level access pattern override (escape hatch).

Replaces shape and strides with an explicit [[stride, count], ...] pattern addressing the underlying storage directly. Analogous to torch.as_strided.

Each stride and count slot may be a Python int or dynamic value only known at runtime — nki.isa.VirtualRegister. Note that when any slot is dynamic, scalar_offset / vector_offset / non-zero indirect_dim are not supported.

sb = nl.ndarray((128, 32), dtype=nl.float32, buffer=nl.sbuf)
# Explicit 2D pattern: partition stride=32, free stride=1
sb.ap(pattern=[[32, 128], [1, 32]])

# With indirect access and dtype reinterpret:
sb.ap(pattern=[[64, 128], [1, 16]], dtype=nl.bfloat16,
      scalar_offset=idx, indirect_dim=1)

# Register-bound stride (dynamic access pattern):
stride0 = nisa.register_alloc(...)
count1 = nisa.register_alloc(...)
sb.ap(pattern=[[stride0, 128], [1, count1]])

Parameters:

pattern – list of [stride, count] pairs defining the access pattern. Each entry may be int or VirtualRegister.
offset – element offset added to the view’s base storage offset. When None (default), inherits the current view’s storage offset unchanged. Pass an explicit integer to compose with the base offset (e.g. offset=0 keeps the base, offset=N shifts by N additional elements).
scalar_offset – dynamic scalar index for indirect access, as a 1x1 index tensor or an NkiScalar. Must be 32-bit when self is in SBUF or PSUM; a 64-bit offset is only valid on an HBM operand, where it reaches past 4GB
vector_offset – per-partition index tensor for indirect access
indirect_dim – dimension in self.shape whose stride scales the indirect scalar/vector offset (default 0)
dtype – reinterpret storage as this dtype (default: tensor’s dtype)

Returns:

new NkiTensor with the explicit access pattern

broadcast(dim: int, size: int) → NkiTensor[source]#

Expand a size-1 dimension to size by repeating elements.

The dimension must have size 1 before broadcasting. No data is copied.

t = nl.ndarray((128, 1, 64), dtype=nl.float32, buffer=nl.sbuf)
t.broadcast(1, 8)  # shape becomes (128, 8, 64)

Constraints.

t.shape[dim] must be 1
On-chip tensors: dim must not be 0 (partition dim)

Parameters:

dim – dimension to broadcast
size – new size for the dimension

Returns:

new NkiTensor view with the broadcasted dimension

buffer: nki.language.MemoryRegion | str = Ellipsis#: The memory region the storage lives in. One of nl.sbuf, nl.psum, nl.hbm, nl.private_hbm, or nl.shared_hbm.

dtype: str = Ellipsis#: Element type, e.g. nl.float32 or nl.bfloat16.

expand_dim(dim: int) → NkiTensor[source]#

Insert a new dimension of size 1 at position dim.

t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t.expand_dim(1)  # shape becomes (128, 1, 64)

Constraints.

On-chip tensors: dim must not be 0
After vector_select: dim must not be 0 (cannot insert before the indirect partition dim; see is_indirect())

Parameters:: dim – position at which to insert the new dimension
Returns:: new NkiTensor view with an additional size-1 dimension

flatten_dims(start_dim: int, end_dim: int) → NkiTensor[source]#

Merge a contiguous range of dimensions into one.

Dimensions start_dim through end_dim (inclusive) are merged into a single dimension. The dimensions must already be contiguous in memory (no permute or non-contiguous slicing across them) so the merged view is itself a valid view of storage.

t = nl.ndarray((128, 2, 3, 4), dtype=nl.float32, buffer=nl.sbuf)
t.flatten_dims(1, 2)  # shape becomes (128, 6, 4)

Constraints.

Dimensions start_dim..end_dim must be contiguous in memory
On-chip tensors: start_dim must be > 0
After vector_select: start_dim must be > 0

Parameters:

start_dim – first dimension to merge (inclusive)
end_dim – last dimension to merge (inclusive)

Returns:

new NkiTensor view with the merged dimension

get_pattern() → list[list[int | VirtualRegister]][source]#

Return the view’s access pattern as [[stride, count], ...].

Useful as a starting point when constructing a new .ap() that keeps most of the current layout intact. The returned pattern pairs each of the view’s dimensions with its current stride, in the same order as shape / strides.

indirect(index: NkiTensor, num_elem: Optional[int] = None) → NkiTensor[source]#

Create an indirect tensor view for Tensor Indirection (TI).

Available on NeuronCore-v4 (Trn3) and later.

TI allows reading or writing a column of data at given free-dimension offsets across contiguous partition dimensions. Can be used as input (gather) or output (scatter) in nisa operations.

Offsets are stored in a snake pattern across partition groups: offset i comes from index[i % G, i // G] where G is the group size (16 for vector/scalar/gpsimd engines, 32 for tensor engine).

Parameters:

index – SBUF tensor containing free-dimension offsets, shape (P, K) where P == self.shape[0].
num_elem – number of offsets to use. Defaults to index.size.

Returns:

new NkiTensor view with TI attached. Output shape is (P, num_elem).

is_contiguous() → bool[source]#

Return True if the view covers storage contiguously (row-major order).

Computed from the current strides: each non-size-1 dimension’s stride must equal the product of the shape sizes of inner dimensions.

For on-chip tensors (SBUF, PSUM), the partition dim (dim 0) is skipped: partitions are physically independent memory banks, so the partition stride does not represent physical contiguity. Contiguity is evaluated per-partition over the free dims only.

is_indirect() → bool[source]#

Return True if this view already uses indirect addressing.

Indirect addressing is produced by dynamic select(), vector_select(), or ap() with scalar_offset / vector_offset. Indirect views cannot be re-indirected, and the dimension that participates in the indirection cannot be further sliced or selected — use this query to guard against those chains.

name: str = Ellipsis#: Optional debug label, used in compiler diagnostics and the profiler.

ndim: int = Ellipsis#: Number of dimensions (len(shape)).

offset: int = Ellipsis#: Element offset (not byte offset) from the start of the underlying storage at which this view begins.

permute(dims: tuple[int, ...]) → NkiTensor[source]#

Reorder tensor dimensions.

Returns a new view with dimensions rearranged according to dims. No data is copied.

t = nl.ndarray((128, 4, 8), dtype=nl.float32, buffer=nl.sbuf)
t.permute((0, 2, 1))  # shape becomes (128, 8, 4)

Constraints.

dims must be a permutation of range(t.ndim)
On-chip tensors: dims[0] must be 0 (partition dim stays outermost)
After vector_select: dims[0] must be 0 (indirect partition dim stays outermost; see is_indirect())

Parameters:: dims – tuple of dimension indices in the desired order
Returns:: new NkiTensor view with reordered dimensions

rearrange(src_pattern: tuple, dst_pattern: tuple, fixed_sizes: Optional[dict[str, int]] = None) → NkiTensor[source]#

Rearrange tensor dimensions using einops-style patterns.

Combines splitting, reordering, and merging dimensions into a single named operation. Patterns are tuples of strings (dimension names) or tuples of strings (grouped dimensions that are split or merged).

t = nl.ndarray((128, 24), dtype=nl.float32, buffer=nl.sbuf)
# Split dim 1 into (h, w), then reorder to (b, w, h):
t.rearrange(('b', ('h', 'w')), ('b', 'w', 'h'), {'h': 4})
# Result shape: (128, 6, 4)

Parameters:

src_pattern – source dimension pattern (tuple of str or tuple-of-str)
dst_pattern – destination dimension pattern (same dimension names)
fixed_sizes – dict mapping dimension names to known sizes (for -1 inference)

Returns:

new NkiTensor view with rearranged dimensions

reinterpret_cast(dtype: str) → NkiTensor[source]#

Alias for view().

Parameters:: dtype – target NKI dtype to reinterpret as
Returns:: new NkiTensor view with the adjusted dtype and shape

reshape(shape: tuple[int, ...]) → NkiTensor[source]#

Reshape the tensor to a new shape without copying data.

The total number of elements must remain the same. Fails if the current memory layout is incompatible with the requested shape (e.g. after a non-contiguous slice or permute).

t = nl.ndarray((128, 4, 6), dtype=nl.float32, buffer=nl.sbuf)
t.reshape((128, 24))       # merge last two dims
t.reshape((128, 2, 12))    # split differently

Constraints.

prod(shape) == prod(t.shape)
On-chip tensors: shape[0] must equal t.shape[0] (partition dim preserved)
After vector_select: shape[0] must equal t.shape[0] (indirect partition dim preserved)
Fails if the current layout is incompatible with the requested shape

Parameters:: shape – tuple of new dimension sizes
Returns:: new NkiTensor view with the requested shape

reshape_dim(dim: int, shape: tuple[int, ...]) → NkiTensor[source]#

Split a single dimension into multiple dimensions.

The product of shape must equal t.shape[dim]. One element of shape may be -1, in which case its value is inferred.

t = nl.ndarray((128, 24), dtype=nl.float32, buffer=nl.sbuf)
t.reshape_dim(1, (4, 6))   # shape becomes (128, 4, 6)
t.reshape_dim(1, (4, -1))  # same result, 6 is inferred

Constraints.

prod(shape) == t.shape[dim]
On-chip tensors: dim must not be 0 (unless shape is trivial, e.g., (128,))
After vector_select: dim must not be 0

Parameters:

dim – dimension to split
shape – tuple of sizes for the new dimensions (may contain one -1)

Returns:

new NkiTensor view with the dimension split

select(dim: int, index: Union[int, NkiTensor, VirtualRegister, NkiScalar]) → NkiTensor[source]#

Select a single element along a dimension, removing it.

The index is either a compile-time int (equivalent to t[:, index, :] when dim=1) or a runtime value resolved on device: an NkiTensor 1-element tile, a VirtualRegister, or an NkiScalar.

t = nl.ndarray((128, 8, 64), dtype=nl.float32, buffer=nl.sbuf)
t.select(1, 3)          # static: shape becomes (128, 64)

# Dynamic select (HBM tensor, index resolved at runtime):
idx = nl.ndarray((1, 1), dtype=nl.uint32, buffer=nl.sbuf)
hbm_t = nl.ndarray((4, 128, 8), dtype=nl.float32, buffer=nl.shared_hbm)
hbm_t.select(0, idx)    # shape becomes (128, 8)

# Same thing driven by a register:
reg = nisa.register_alloc()
nisa.register_load(reg, idx)
hbm_t.select(0, reg)

Constraints.

Static: 0 <= index < t.shape[dim]; runtime indices are not bounds-checked, and out-of-range is undefined behaviour
One runtime index per view (no chaining) — see is_indirect()
Runtime on-chip: dim must not be 0 (partition dim)
Static selection cannot target a dimension already indirected

Parameters:

dim – dimension to select from
index – compile-time int, or a runtime NkiTensor / VirtualRegister / NkiScalar

Returns:

new NkiTensor view with the dimension removed

shape: tuple[int, ...] = Ellipsis#: Tuple giving the count of each dimension of the view. Each element is an integer, or a VirtualRegister for dynamic-access views constructed via ap().

size: int = Ellipsis#: Total number of elements (prod(shape)). None when any element of shape is a VirtualRegister.

slice(dim: int, start: int, end: int, step: int = 1) → NkiTensor[source]#

Slice along a single dimension.

Returns a view selecting elements from start to end (exclusive) with the given step. Equivalent to t[:, start:end:step, :] when dim=1.

t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t.slice(1, 8, 24, 1)   # shape becomes (128, 16)
t.slice(1, 0, 64, 2)   # shape becomes (128, 32)

Constraints.

0 <= start < end <= t.shape[dim]
step >= 1
On an indirect view (see is_indirect()), cannot slice a dimension that participates in the indirection.

Parameters:

dim – dimension to slice
start – start index (inclusive)
end – end index (exclusive)
step – step size (default 1)

Returns:

new NkiTensor view with the sliced dimension

squeeze_dim(dim: int) → NkiTensor[source]#

Remove a dimension of size 1.

t = nl.ndarray((128, 1, 64), dtype=nl.float32, buffer=nl.sbuf)
t.squeeze_dim(1)  # shape becomes (128, 64)

Constraints.

t.shape[dim] must be 1
On-chip tensors: dim must not be 0
After vector_select: dim must not be 0

Parameters:: dim – dimension to remove (must have size 1)
Returns:: new NkiTensor view with the dimension removed

strides: tuple[int, ...] = Ellipsis#: Tuple with stride per dimension. Each stride is the element step between consecutive indices along that dimension in the underlying storage. Each entry is an integer, or a VirtualRegister for dynamic-access views constructed via ap(). A stride of 0 means the dimension is broadcast.

vector_select(dim: int, vector_offset: NkiTensor) → NkiTensor[source]#

Per-partition indirect addressing using a vector of offsets.

Each partition uses its own index from vector_offset as the base address for the selected dimension. The output partition count is vector_offset.shape[0].

This is used for gather-style operations where different partitions read from different locations in the source tensor.

hbm_t = nl.ndarray((64, 128, 8), dtype=nl.float32, buffer=nl.shared_hbm)
offsets = nl.ndarray((128, 1), dtype=nl.int32, buffer=nl.sbuf)
# Each of 128 partitions reads from a different row of hbm_t
hbm_t.vector_select(0, offsets)  # shape becomes (128, 128, 8)

Constraints.

dim must be 0
Only supported on HBM tensors
Only one dynamic select per tensor (cannot combine with a prior dynamic select() or vector_select()); check is_indirect() to guard
The result is an indirect view: the selected dimension cannot be further sliced or selected.

Parameters:

dim – dimension to apply indirect addressing (must be 0)
vector_offset – SBUF tensor with per-partition indices, shape (num_partitions, 1)

Returns:

new NkiTensor view with dim 0 size set to vector_offset.shape[0]

view(dtype: str) → NkiTensor[source]#

Reinterpret the tensor’s data as a different dtype.

No data is copied. Equivalent to numpy.ndarray.view(dtype) or C++ reinterpret_cast. The last dimension’s size is scaled by the ratio of dtype sizes: reinterpreting a float32 tensor as uint8 multiplies the last-dim count by 4; reinterpreting uint8 as float32 divides it by 4.

t = nl.ndarray((128, 64), dtype=nl.float32, buffer=nl.sbuf)
t.view(nl.uint8)     # shape becomes (128, 256), 4x expansion
t.view(nl.int32)     # shape stays (128, 64), same-size reinterpret

u = nl.ndarray((128, 256), dtype=nl.uint8, buffer=nl.sbuf)
u.view(nl.float32)   # shape becomes (128, 64), 4x contraction

Constraints.

The last dimension must be contiguous in memory
For contraction (larger dtype): last-dim size must be divisible by the ratio
Not supported after dynamic / vector select

Parameters:: dtype – target NKI dtype to reinterpret as
Returns:: new NkiTensor view with the adjusted dtype and shape

This document is relevant for: Trn2, Trn3

nki.language.NkiTensor

Contents

nki.language.NkiTensor#