viva_tensor/core/shape

Shape manipulation - reshape, slice, concat, stack.

The “plumbing” of tensor operations. Not glamorous, but essential.

Philosophy: these ops should be zero-copy when possible (via strides), but we’re not there yet. Current implementation copies data. TODO: implement as views where safe (reshape of contiguous tensor, slice, etc.)

Reshape is particularly tricky because it changes how we interpret memory. A [2,3] tensor reshaped to [3,2] has the same data but different semantics. This only works if the tensor is contiguous (strides match row-major order).

Values

pub fn concat(tensors: List(tensor.Tensor)) -> tensor.Tensor

Concat 1D tensors. Just appends the data, nothing fancy.

pub fn concat_axis(
  tensors: List(tensor.Tensor),
  axis: Int,
) -> Result(tensor.Tensor, error.TensorError)

Concat along arbitrary axis.

This function is gnarly. The general case requires computing which source tensor each output index maps to, then translating coordinates. O(n) where n is total output size, but the constant factor is high due to all the index arithmetic. For axis=0, we fast-path to simple concatenation.

pub fn expand_dims(t: tensor.Tensor, axis: Int) -> tensor.Tensor

Alias for unsqueeze - expand dimensions

pub fn flatten(t: tensor.Tensor) -> tensor.Tensor

Flatten to 1D tensor

pub fn reshape(
  t: tensor.Tensor,
  new_shape: List(Int),
) -> Result(tensor.Tensor, error.TensorError)

Reshape to new dimensions. Total size must match (obviously).

pub fn slice(
  t: tensor.Tensor,
  start: List(Int),
  lengths: List(Int),
) -> Result(tensor.Tensor, error.TensorError)

General slice - specify start indices and lengths for each dimension. This one’s tricky for n-dimensional tensors.

pub fn squeeze(t: tensor.Tensor) -> tensor.Tensor

Remove all dimensions of size 1

pub fn squeeze_axis(
  t: tensor.Tensor,
  axis: Int,
) -> Result(tensor.Tensor, error.TensorError)

Remove dimension at specific axis if it’s 1

pub fn stack(
  tensors: List(tensor.Tensor),
  axis: Int,
) -> Result(tensor.Tensor, error.TensorError)

Stack tensors along a new axis

pub fn take_first(t: tensor.Tensor, n: Int) -> tensor.Tensor

Take first n elements (along axis 0).

pub fn take_last(t: tensor.Tensor, n: Int) -> tensor.Tensor

Take last N elements along first axis

pub fn unsqueeze(t: tensor.Tensor, axis: Int) -> tensor.Tensor

Add dimension of size 1 at specified axis

Search Document