Tensor is a memoryview-backed object with NumPy-like metadata.
It is the central container for strided views, transpose, reshape, flatten, and axis reductions.
import numpy as np
import numkong as nk
t = nk.Tensor(np.arange(12, dtype=np.float32).reshape(3, 4))
print(t.shape, t.dtype, t.ndim, t.strides, t.itemsize, t.nbytes)
print(np.asarray(t)) # zero-copy array view when layout allows it
print(t.T.shape) # transposed Tensor view
print(t.reshape(2, 6).shape)
print(t.flatten().shape)
# Slicing — row, column, and scalar access
row0 = t[0, :] # first row, shape (4,)
col2 = t[:, 2] # third column, strided view, shape (3,)
val = t[1, 2] # scalar element access → 6.0
# Reductions compose with sliced views
idx = col2.argmin() # index of the minimum in the third column
mn, i0, mx, i1 = col2.minmax()The important layout rules are:
Tensor preserves shape and byte strides.out buffers.| API family | Input requirement | Output requirement |
|---|---|---|
Dense distances (dot, euclidean, etc.) | Rows must be contiguous (strides[last] <= itemsize). Strided rows (sliced columns) are rejected. | out= can have any stride along dim 0, but inner dim must be contiguous. |
cdist | Same as dense distances | out= must be rank-2 with shape (a.count, b.count) |
Elementwise (scale, blend, fma) | Arbitrary strides (strided views are supported) | out= must match input shape; strides are preserved |
Packed matrix (dots_packed) | Left operand: rank-2, contiguous rows, no negative strides | Output: C-contiguous with expected dtype |
Symmetric (dots_symmetric) | Contiguous rows | out=: C-contiguous square matrix |
Tensor reductions (sum, min, argmin, etc.) | Arbitrary strides (strided views supported) | N/A (returns scalar or reduced tensor) |
Builds configured with -C cmake.define.NUMKONG_CUDA_ARCH=all or -C cmake.define.NUMKONG_ROCM_ARCH=all also carry the CUDA or ROCm kernels. nk.cuda_count_devices() counts the devices this process sees, zero without the runtime or its kernels, and nk.cuda_capabilities_enabled(0) reports the mask device 0 dispatches with, raising ValueError past the last. nk.cuda_capabilities_detected and nk.cuda_capabilities_compiled answer as their CPU twins do, nk.rocm_* and nk.metal_* do the same for the other vendors, and cpu_configure_thread is CPU only. nk.from_dlpack imports a CUDA or ROCm tensor without a copy and records its device.
Only dots_pack, the dots, angulars and euclideans _packed and _symmetric functions, and attention_pack, attention_packed and attention_packed_gradients run on it, with that device’s enabled capabilities.
The shape functions also take it for their outputs, running the mask their packs were copied with.
Element access, NumPy conversion, cdist, maxsim and every other CPU kernel raise BufferError rather than read device memory.
No verb allocates, so every GPU call writes into an out= its device reaches.
The dots and attention families take that as a Tensor on the same device, and mixing devices raises ValueError.
The shape functions take any memory the device reaches, a Tensor or a buffer over managed or pinned memory, which the runtime checks per call, raising BufferError for the rest.
Unified memory comes only from an explicit copy, as in nk.shapes_pack(clouds).copy(capabilities=gpu) and nk.zeros(shape).copy(capabilities=gpu), which frees with the tensor and its views and refuses resize() and reserve().
A GPU call queues one launch on stream=, an integer handle defaulting to 0, the legacy stream, and returns before it finishes, as PyTorch does, and nk.synchronize(capabilities, stream=stream) waits for it. nk.stream_init_cuda(ordinal) makes such a handle on any device by its ordinal, and nk.stream_free_cuda(stream) frees it once synchronized, with rocm_ and metal_ twins.
Metal tensors are host-readable and keep running on the CPU kernels, except the shape functions, which run on Metal over blocks of its allocator.
import torch
import numkong as nk
left = torch.randn(4096, 768, device="cuda", dtype=torch.bfloat16)
right = torch.randn(8192, 768, device="cuda", dtype=torch.bfloat16)
capabilities = nk.cuda_capabilities_enabled(right.device.index)
packed_storage = torch.empty(nk.PackedMatrix.pack_size(8192, 768, "bf16", capabilities=capabilities),
device="cuda", dtype=torch.uint8)
result = torch.empty(4096, 8192, device="cuda", dtype=torch.float32)
stream = torch.cuda.current_stream().cuda_stream
packed = nk.dots_pack(nk.from_dlpack(right), out=nk.from_dlpack(packed_storage), stream=stream)
nk.dots_packed(nk.from_dlpack(left), packed, out=nk.from_dlpack(result), stream=stream)
nk.synchronize(capabilities, stream=stream)TensorN-dimensional tensor type returned by NumKong operations.44class Tensor(memoryview)Supports NumPy-like properties and buffer protocol for interoperability.
Packed dtypes 'uint1', 'int4', 'uint4', and 'e2m1' count logical dimensions in shape;
the last dimension must be a multiple of the values per byte, and the buffer protocol,
__array_interface__, nbytes, and DLPack expose whole bytes.
TensorConstruct a Tensor by copying data from a buffer-protocol object; packed dtypes read its bytes.Tensor(
array_like: _BufferType,
/,
*,
dtype: str | _MiniFloatType | None = None,
)__len__Return the length of the first dimension.def __len__() -> int__iter__Iterate over the first dimension.def __iter__() -> Any__getitem__Get an element or sub-tensor by index or slice.def __getitem__(
key: int | slice | tuple[int | slice, ...],
) -> float | Tensor__repr__*__repr__× 2Return a string representation.def __repr____repr__() -> str__*__× 2Convert 0D tensor to Python float.__str__Return a pretty-printed string representation.def __str__() -> str__eq__Compare tensors for equality.def __eq__(other: Any) -> bool__ne__Compare tensors for inequality.def __ne__(other: Any) -> bool__add__Element-wise sum with another same-shape Tensor, or a broadcast scalar.def __add__(other: Tensor | float | int) -> Tensor__sub__Element-wise difference with another same-shape Tensor, or a broadcast scalar.def __sub__(other: Tensor | float | int) -> Tensor__mul__Element-wise product with another same-shape Tensor, or a broadcast scalar.def __mul__(other: Tensor | float | int) -> Tensor__matmul__Dot-product matrix multiplication against a pre-packed matrix; see dots_packed.def __matmul__(other: PackedMatrix) -> Tensor__neg__Return the element-wise negation.def __neg__() -> Tensor__pos__Return a copy of the tensor.def __pos__() -> Tensor__dlpack__Return a DLPack capsule for zero-copy exchange with PyTorch, JAX, NumPy, etc.def __dlpack__(
*,
stream: Any | None = None,
max_version: tuple[int, int] | None = None,
dl_device: tuple[int, int] | None = None,
copy: bool | None = None,
) -> Any__dlpack_device__Return the DLPack device tuple.__array__Convert to a NumPy array, sharing data when possible.def __array__(
dtype: str | None = None,
*,
copy: bool | None = None,
) -> AnycopyReturn a deep copy of the tensor, or write it into out and return out.def copy(
*,
out: Tensor | None = None,
capabilities: Capability | int | None = None,
stream: int | None = None,
) -> TensorreshapeReturn tensor reshaped to given dimensions.def reshape(*shape: int = ()) -> TensorflattenReturn a flattened 1D view, or a copy if non-contiguous or out is given.def flatten(*, out: Tensor | None = None) -> TensorsqueezeRemove the given size-1 dimensions, or every size-1 dimension when axes is empty.def squeeze(*axes: int = ()) -> TensorastypeReturn a copy cast to the given dtype, quantizing to a ScaledTensor for a block-scaled target.def astype(
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType | _BlockScaledTypeName,
/,
*,
out: Tensor | None = None,
) -> Tensor | ScaledTensorresizeReshape in place within capacity without moving storage; returns self.def resize(*shape: int = ()) -> TensorreserveGrow the allocated capacity to at least capacity elements (may move); preserves contents.def reserve(capacity: int) -> NoneclearReset to an empty shape (size 0) while keeping the allocated capacity.def clear() -> NonesumReturn the sum of all elements.def sum(
axis: int | tuple[int, ...] | None = None,
*,
keepdims: bool = False,
out: Tensor | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> float | int | TensorcapabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
|---|---|---|
streamint | None | A GPU stream pointer, left None on the CPU. |
normReturn the L2 norm.def norm(
axis: int | tuple[int, ...] | None = None,
*,
keepdims: bool = False,
out: Tensor | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> float | TensorcapabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
|---|---|---|
streamint | None | A GPU stream pointer, left None on the CPU. |
minReturn the minimum element, or None if all elements are NaN.def min(
axis: int | tuple[int, ...] | None = None,
*,
keepdims: bool = False,
out: Tensor | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> float | int | Tensor | NonecapabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
|---|---|---|
streamint | None | A GPU stream pointer, left None on the CPU. |
maxReturn the maximum element, or None if all elements are NaN.def max(
axis: int | tuple[int, ...] | None = None,
*,
keepdims: bool = False,
out: Tensor | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> float | int | Tensor | NonecapabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
|---|---|---|
streamint | None | A GPU stream pointer, left None on the CPU. |
argminReturn the index of the minimum element, or None if all elements are NaN.def argmin(
axis: int | None = None,
*,
keepdims: bool = False,
out: Tensor | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> int | Tensor | NonecapabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
|---|---|---|
streamint | None | A GPU stream pointer, left None on the CPU. |
argmaxReturn the index of the maximum element, or None if all elements are NaN.def argmax(
axis: int | None = None,
*,
keepdims: bool = False,
out: Tensor | None = None,
capabilities: Capability | int | None = ...,
stream: int | None = ...,
) -> int | Tensor | NonecapabilitiesCapability | int | None | Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them. | |
|---|---|---|
streamint | None | A GPU stream pointer, left None on the CPU. |
shapeShape of the tensor as a tuple of logical dimensions.shape: tuple[int, ...]dtypeData type of the tensor elements (e.g., 'float64').dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeNamendimNumber of dimensions.ndim: intsizeTotal number of elements.size: intnbytesTotal number of bytes.nbytes: intstridesStrides of the tensor in bytes.strides: tuple[int, ...]itemsizeSize of each element in bytes.itemsize: int__array_interface__NumPy array interface dict for legacy interoperability.__array_interface__: dict[str, Any]TTranspose of the tensor.is_contiguousWhether the tensor is C-contiguous in memory.is_contiguous: booldata_ptrInteger address of the underlying data buffer.data_ptr: intcapacityAllocated element capacity of the owned buffer (>= size).capacity: intScaledTensorBlock-scaled tensor: the OCP MX family plus NVIDIA NVFP4.10class ScaledTensorHolds packed elements, per-block scales, and an optional per-tensor scale.
Produced by Tensor.astype('nvfp4'), Tensor.astype('mxfp4'), and similar block-scaled targets.
astypeMaterialize to a dense dtype, or transcode to another block-scaled dtype.def astype(
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType | _BlockScaledTypeName,
/,
) -> Tensor | ScaledTensorresizeResize the packed elements and block scales in lockstep, within capacity; returns self.def resize(*shape: int = ()) -> ScaledTensor__getitem__Slice elements and block_scales in lockstep; the last axis must stay block-aligned.def __getitem__(
key: int | slice | tuple[int | slice, ...],
) -> ScaledTensorelementsPacked sub-byte element Tensor, DLPack-exportable.elements: Tensorblock_scalesPer-block scale Tensor, DLPack-exportable.block_scales: Tensortensor_scalePer-tensor float multiplier, None for MX formats.tensor_scale: float | Noneblock_sizeNumber of elements per block.block_size: intdtypeComposite block-scaled dtype name.dtype: _BlockScaledTypeNameshapeLogical, dense shape of the quantized tensor.shape: tuple[int, ...]capacityAllocated element capacity of the packed elements buffer.capacity: intfrom_pointerCreate a zero-copy Tensor view over memory at a raw address.def from_pointer(
address: int,
shape: int | tuple[int, ...],
dtype: _DTypeLike,
/,
*,
strides: tuple[int, ...] | None = None,
owner: Any = None,
) -> Tensoraddressint | Integer address of the data buffer. | |
|---|---|---|
shapeint | tuple[int, ...] | Shape of the view, an int or a tuple of ints. | |
dtype_DTypeLike | Element dtype, like 'float32' or 'bf16'. | |
stridestuple[int, ...] | None | Byte strides, C-contiguous when None. | |
ownerAny | Object the view keeps alive, so the memory it owns is not collected. |
from_dlpackZero-copy import a tensor from any DLPack-compatible producer.def from_dlpack(obj: Any, /) -> TensorAccepts a PyCapsule named 'dltensor' or 'dltensor_versioned', or any object
with a __dlpack__ method (e.g. torch.Tensor, numpy.ndarray,
jax.Array). Only CPU tensors are supported.
emptyCreate an uninitialized Tensor of the given shape and dtype.def empty(
shape: int | tuple[int, ...],
/,
*,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType = 'float32',
) -> Tensor*× 2Create a Tensor of the given shape and dtype filled with zeros.def zeros(
shape: int | tuple[int, ...],
/,
*,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType = 'float32',
) -> Tensor| Symbol |
|---|
zeros |
ones |
fullCreate a Tensor of the given shape and dtype filled with fill_value.def full(
shape: int | tuple[int, ...],
fill_value: int | float,
/,
*,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType = 'float32',
) -> TensoriotaTensor filled with seed, seed+1, seed+2, ... reshaped to shape.def iota(
shape: int | tuple[int, ...],
/,
seed: int | float = 0,
*,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType = 'float32',
) -> TensordiagonalSquare n x n tensor with seed on the main diagonal and zeros elsewhere.def diagonal(
n: int,
/,
seed: int | float = 1,
*,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType = 'float32',
) -> TensorhashTensor filled with deterministic pseudo-random bits (shape, dtype, seed → reproducible bytes).def hash(
shape: int | tuple[int, ...],
/,
seed: int = 0,
*,
dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _ComplexTypeName | _MiniFloatType = 'float32',
) -> TensorastypeCast an N-D buffer to dtype without first copying it into a Tensor.def astype(
a: _BufferType,
dtype: _DTypeLike,
/,
*,
out: _BufferType | None = None,
) -> TensorA supplied out must be writable, C-contiguous, have the exact input
shape and requested dtype, and must not overlap a. Returns out when
provided, otherwise a new Tensor.