Tensor

Multi-dimensional array with hardware acceleration support

Overview

The Tensor class is the fundamental data structure in Neurenix, similar to tensors in PyTorch or TensorFlow. It provides a multi-dimensional array with automatic device management and gradient tracking.

Class Definition

class Tensor:
    def __init__(
        self,
        data: Union[np.ndarray, List, Tuple, Tensor, None] = None,
        shape: Optional[Sequence[int]] = None,
        dtype: Optional[Union[DType, str]] = None,
        device: Optional[Device] = None,
        requires_grad: bool = False,
    )

Parameters

data · Union[np.ndarray, List, Tuple, Tensor, None]

The data to initialize the tensor with. Can be a NumPy array, a list, a tuple, another Tensor, or None (for uninitialized tensor).

shape · Optional[Sequence[int]]

The shape of the tensor. If None, inferred from data.

dtype · Optional[Union[DType, str]]

The data type of the tensor. If None, inferred from data. Options: ‘float32’, ‘float64’, ‘int32’, ‘int64’, ‘bool’.

device · Optional[Device]

The device to store the tensor on. If None, uses the default device.

requires_grad · bool · default: "False"

Whether to track gradients for this tensor.

Properties

shape

@property
def shape(self) -> Tuple[int, ...]

Get the shape of the tensor.

return · Tuple[int, ...]

The shape of the tensor as a tuple of integers.

dtype

@property
def dtype(self) -> DType

Get the data type of the tensor.

device

@property
def device(self) -> Device

Get the device where the tensor is stored.

requires_grad

@property
def requires_grad(self) -> bool

Check if the tensor requires gradients.

grad

@property
def grad(self) -> Optional[Tensor]

Get the gradient of the tensor.

Methods

numpy

def numpy(self) -> np.ndarray

Convert the tensor to a NumPy array. This operation will copy the tensor data from the device to the CPU if necessary.

return · np.ndarray

A NumPy array with the tensor data.

to

def to(self, device: Device, non_blocking: bool = False) -> Tensor

Move the tensor to the specified device.

device · Device

The target device.

non_blocking · bool · default: "False"

If True, the copy will be performed asynchronously. Only has an effect for CUDA/ROCm devices.

return · Tensor

A new tensor on the target device.

hot_swap_device

def hot_swap_device(self, device: Device) -> None

Hot-swap the tensor to a different device without creating a new tensor. This method changes the device of the tensor in-place.

reshape

def reshape(self, *shape) -> Tensor

Reshape the tensor to the given shape.

shape · *args

The new shape of the tensor.

return · Tensor

A new tensor with the given shape.

transpose

def transpose(self, dim0: int = 0, dim1: int = 1) -> Tensor

Transpose the tensor along the given dimensions.

matmul

def matmul(self, other: Tensor) -> Tensor

Matrix multiplication with another tensor.

mean

def mean(self, dim: Optional[int] = None, keepdim: bool = False) -> Tensor

Compute the mean along the specified dimension.

sum

def sum(self, dim: Optional[int] = None, keepdim: bool = False) -> Tensor

Compute the sum along the specified dimension.

clone

def clone(self) -> Tensor

Create a clone of this tensor.

backward

def backward(self)

Compute gradients through the computation graph.

Activation Functions

relu

def relu(self, inplace: bool = False) -> Tensor

Apply the rectified linear unit function element-wise.

sigmoid

def sigmoid(self) -> Tensor

Apply the sigmoid function element-wise.

tanh

def tanh(self) -> Tensor

Apply the hyperbolic tangent function element-wise.

softmax

def softmax(self, dim: int = -1) -> Tensor

Apply the softmax function along the specified dimension.

log_softmax

def log_softmax(self, dim: int = -1) -> Tensor

Apply the log softmax function along the specified dimension.

leaky_relu

def leaky_relu(self, negative_slope: float = 0.01, inplace: bool = False) -> Tensor

Apply the leaky rectified linear unit function element-wise.

gelu

def gelu(self, approximate: bool = False) -> Tensor

Apply the Gaussian error linear unit function element-wise.

Static Methods

zeros

@staticmethod
def zeros(
    shape: Sequence[int],
    dtype: Optional[DType] = None,
    device: Optional[Device] = None,
    requires_grad: bool = False
) -> Tensor

Create a tensor filled with zeros.

ones

@staticmethod
def ones(
    shape: Sequence[int],
    dtype: Optional[DType] = None,
    device: Optional[Device] = None,
    requires_grad: bool = False
) -> Tensor

Create a tensor filled with ones.

randn

@staticmethod
def randn(
    shape: Sequence[int],
    dtype: Optional[DType] = None,
    device: Optional[Device] = None,
    requires_grad: bool = False
) -> Tensor

Create a tensor filled with random numbers from a normal distribution.

stack

@staticmethod
def stack(tensors: List[Tensor], dim: int = 0) -> Tensor

Stack tensors along a new dimension.

cat

@staticmethod
def cat(tensors: List[Tensor], dim: int = 0) -> Tensor

Concatenate tensors along an existing dimension.

no_grad

@staticmethod
@contextmanager
def no_grad()

Context manager to disable gradient computation.

Example Usage

import neurenix as nx

# Create a tensor from a list
x = nx.Tensor([[1, 2, 3], [4, 5, 6]])
print(x.shape)  # (2, 3)

# Create a tensor with zeros
zeros = nx.Tensor.zeros((3, 4), dtype=nx.DType.FLOAT32)

# Create a random tensor
rand = nx.Tensor.randn((2, 3, 4))

# Reshape a tensor
reshaped = x.reshape(3, 2)

# Apply activation functions
activated = x.relu()
softmax_output = x.softmax(dim=1)

# Matrix multiplication
y = nx.Tensor.randn((3, 5))
result = x.matmul(y)

# Move to device
device = nx.Device(nx.DeviceType.CUDA, 0)
x_cuda = x.to(device)

# Gradient tracking
x = nx.Tensor([1.0, 2.0, 3.0], requires_grad=True)
y = x * 2
z = y.sum()
z.backward()
print(x.grad)  # Gradients

# Context manager for inference
with nx.Tensor.no_grad():
    prediction = model(input_tensor)