Mojo struct

Tensor

A tensor type which owns its underlying data and is parameterized on DType.

Example:

from max.tensor import Tensor, TensorSpec, TensorShape
from utils.index import Index

def main():
    height = 256
    width = 256
    channels = 3

    # Create the tensor of dimensions height, width, channels
    # and fill with random values.
    image = Tensor[DType.float32].rand(TensorShape(height, width, channels))

    # Declare the grayscale image.
    spec = TensorSpec(DType.float32, height, width)
    gray_scale_image = Tensor[DType.float32](spec)

    # Perform the RGB to grayscale transform.
    for y in range(height):
        for x in range(width):
            r = image[y, x, 0]
            g = image[y, x, 1]
            b = image[y, x, 2]
            gray_scale_image[Index(y, x)] = 0.299 * r + 0.587 * g + 0.114 * b

    print(gray_scale_image.shape())
from max.tensor import Tensor, TensorSpec, TensorShape
from utils.index import Index

def main():
    height = 256
    width = 256
    channels = 3

    # Create the tensor of dimensions height, width, channels
    # and fill with random values.
    image = Tensor[DType.float32].rand(TensorShape(height, width, channels))

    # Declare the grayscale image.
    spec = TensorSpec(DType.float32, height, width)
    gray_scale_image = Tensor[DType.float32](spec)

    # Perform the RGB to grayscale transform.
    for y in range(height):
        for x in range(width):
            r = image[y, x, 0]
            g = image[y, x, 1]
            b = image[y, x, 2]
            gray_scale_image[Index(y, x)] = 0.299 * r + 0.587 * g + 0.114 * b

    print(gray_scale_image.shape())

Parameters

type (DType): The underlying element type of the tensor.

Implemented traits

AnyType, CollectionElement, Copyable, EqualityComparable, Formattable, Movable, Stringable

Methods

`init`

__init__(inout self: Self)

Default initializer for TensorShape.

__init__(inout self: Self, other: Self)

Creates a deep copy of an existing tensor.

Args:

other (Self): The tensor to copy from.

__init__(inout self: Self, *dims: Int)

Allocates a tensor using the shape provided.

Args:

*dims (Int): The tensor dimensions.

__init__(inout self: Self, owned shape: TensorShape)

Allocates a tensor using the shape provided.

Args:

shape (TensorShape): The tensor shape.

__init__(inout self: Self, owned spec: TensorSpec)

Allocates a tensor using the spec provided.

Args:

spec (TensorSpec): The tensor spec.

__init__(inout self: Self, shape: Tuple[element_types])

Allocates a tensor using the shape provided.

Args:

shape (Tuple[element_types]): The tensor shape.

__init__(inout self: Self, owned shape: TensorShape, owned ptr: UnsafePointer[SIMD[type, 1], 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1])

Initializes a Tensor from the pointer and shape provided. The caller relinquishes the ownership of the pointer being passed in.

Args:

shape (TensorShape): The tensor shapes.
ptr (UnsafePointer[SIMD[type, 1], 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]): The data pointer.

__init__(inout self: Self, owned spec: TensorSpec, owned ptr: UnsafePointer[SIMD[type, 1], 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1])

Initializes a Tensor from the pointer and shape provided. The caller relinquishes the ownership of the pointer being passed in.

Args:

spec (TensorSpec): The tensor spec.
ptr (UnsafePointer[SIMD[type, 1], 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]): The data pointer.

__init__(inout self: Self, shape: TensorShape, *data: SIMD[type, 1])

Initializes a Tensor from the shape and data provided. If a single scalar is passed in, then the scalar is splatted to all elements in the tensor.

Args:

shape (TensorShape): The tensor shape.
*data (SIMD[type, 1]): Elements to place into the created tensor.

__init__(inout self: Self, shape: TensorShape, owned list: List[SIMD[type, 1], hint_trivial_type])

Initializes a 1-dimensional Tensor from the provided list.

Args:

shape (TensorShape): The tensor shape.
list (List[SIMD[type, 1], hint_trivial_type]): The list to construct this Tensor from.

__init__(inout self: Self, owned list: List[SIMD[type, 1], hint_trivial_type])

Initializes a 1-dimensional Tensor from the provided list.

Args:

list (List[SIMD[type, 1], hint_trivial_type]): The list to construct this Tensor from.

`copyinit`

__copyinit__(inout self: Self, other: Self)

Creates a deep copy of an existing tensor.

Args:

other (Self): The tensor to copy from.

`moveinit`

__moveinit__(inout self: Self, owned existing: Self)

Move initializer for the tensor.

Args:

existing (Self): The tensor to move.

`del`

__del__(owned self: Self)

Delete the spec and release any owned memory.

`getitem`

__getitem__(self: Self, index: Int) -> SIMD[type, 1]

Gets the value at the specified index.

Args:

index (Int): The index of the value to retrieve.

Returns:

The value at the specified indices.

__getitem__(self: Self, *indices: Int) -> SIMD[type, 1]

Gets the value at the specified indices.

Args:

*indices (Int): The indices of the value to retrieve.

Returns:

The value at the specified indices.

__getitem__(self: Self, indices: VariadicList[Int]) -> SIMD[type, 1]

Gets the value at the specified indices.

Args:

indices (VariadicList[Int]): The indices of the value to retrieve.

Returns:

The value at the specified indices.

__getitem__[len: Int](self: Self, indices: StaticIntTuple[len]) -> SIMD[type, 1]

Gets the SIMD value at the specified indices.

Parameters:

len (Int): The length of the indecies.

Args:

indices (StaticIntTuple[len]): The indices of the value to retrieve.

Returns:

The value at the specified indices.

`setitem`

__setitem__(inout self: Self, index: Int, val: SIMD[type, 1])

Sets the value at the specified index.

Args:

index (Int): The index of the value to set.
val (SIMD[type, 1]): The value to store.

__setitem__(inout self: Self, indices: VariadicList[Int], val: SIMD[type, 1])

Sets the value at the specified indices.

Args:

indices (VariadicList[Int]): The indices of the value to set.
val (SIMD[type, 1]): The value to store.

__setitem__[len: Int](inout self: Self, indices: StaticIntTuple[len], val: SIMD[type, 1])

Sets the value at the specified indices.

Parameters:

len (Int): The length of the indecies.

Args:

indices (StaticIntTuple[len]): The indices of the value to set.
val (SIMD[type, 1]): The value to store.

`eq`

__eq__(self: Self, other: Self) -> Bool

Returns True if the two tensors are the same and False otherwise.

Args:

other (Self): The other Tensor to compare against.

Returns:

True if the two tensors are the same and False otherwise.

`ne`

__ne__(self: Self, other: Self) -> Bool

Returns True if the two tensors are not the same and False otherwise.

Args:

other (Self): The other Tensor to compare against.

Returns:

True if the two tensors are the not the same and False otherwise.

`add`

__add__(self: Self, other: Self) -> Self

Adds this tensor with another tensor.

Constraints:

The two tensors must have the same rank, type, and dimensions.

Args:

other (Self): The RHS of the add operation.

Returns:

The addition of both tensors.

__add__(self: Self, other: SIMD[type, 1]) -> Self

Adds this tensor with a scalar.

Args:

other (SIMD[type, 1]): The RHS of the add operation.

Returns:

The addition result.

`sub`

__sub__(self: Self, other: Self) -> Self

Subtracts a tensor from this tensor.

Constraints:

The two tensors must have the same rank, type, and dimensions.

Args:

other (Self): The RHS of the sub operation.

Returns:

The addition of both tensors.

__sub__(self: Self, other: SIMD[type, 1]) -> Self

Subtracts a scalar from this tensor.

Args:

other (SIMD[type, 1]): The RHS of the sub operation.

Returns:

The subtraction result.

`mul`

__mul__(self: Self, other: Self) -> Self

Multiplies this tensor with another tensor.

Constraints:

The two tensors must have the same rank, type, and dimensions.

Args:

other (Self): The RHS of the mul operation.

Returns:

The multiplication of both tensors.

__mul__(self: Self, other: SIMD[type, 1]) -> Self

Multiplies this tensor with a scalar.

Args:

other (SIMD[type, 1]): The RHS of the mul operation.

Returns:

The multiplication result.

`truediv`

__truediv__(self: Self, other: Self) -> Self

Divides this tensor by another tensor.

TODO: Change the return type if inputs are int

Constraints:

The two tensors must have the same rank, type, and dimensions.

Args:

other (Self): The RHS of the div operation.

Returns:

The division of both tensors.

__truediv__(self: Self, other: SIMD[type, 1]) -> Self

Divides this tensor by a scalar.

Args:

other (SIMD[type, 1]): The RHS of the div operation.

Returns:

The division result.

`pow`

__pow__(self: Self, exp: Int) -> Self

Returns a copy of the tensor with each element raised to the power of exponent.

Constraints:

For integral values the exponent cannot be negative.

Args:

exp (Int): Integer power to raise tensor to.

Returns:

An exponentiated copy of tensor.

`radd`

__radd__(self: Self, other: SIMD[type, 1]) -> Self

Adds this tensor with a scalar.

Args:

other (SIMD[type, 1]): The LHS of the add operation.

Returns:

The addition result.

`rsub`

__rsub__(self: Self, other: SIMD[type, 1]) -> Self

Subtracts this tensor from a scalar.

Args:

other (SIMD[type, 1]): The LHS of the sub operation.

Returns:

The addition result.

`rmul`

__rmul__(self: Self, other: SIMD[type, 1]) -> Self

Multiplies this tensor with a scalar.

Args:

other (SIMD[type, 1]): The LHS of the mul operation.

Returns:

The multiplication result.

`rtruediv`

__rtruediv__(self: Self, other: SIMD[type, 1]) -> Self

Divides a scalar by this tensor, broadcasting elementwise.

Args:

other (SIMD[type, 1]): The LHS of the div operation.

Returns:

The division result.

`ipow`

__ipow__(inout self: Self, exponent: Int)

In-place pow operator.

Raises each element of the tensor to the power of exponent in place.

Constraints:

For integral values the exponent cannot be negative.

Args:

exponent (Int): Integer power to raise tensor to.

`rand`

static rand(owned shape: TensorShape) -> Self

Constructs a new tensor with the specified shape and fills it with random elements.

Args:

shape (TensorShape): The tensor shape.

Returns:

A new tensor of specified shape and filled with random elements.

`randn`

static randn(owned shape: TensorShape, mean: SIMD[float64, 1] = #kgen.float_literal<0|1>, variance: SIMD[float64, 1] = #kgen.float_literal<1|1>) -> Self

Constructs a new Tensor from the shape and fills it with random values from a Normal(mean, variance) distribution.

Constraints:

The type should be floating point.

Args:

shape (TensorShape): The shape of the Tensor to fill with random values.
mean (SIMD[float64, 1]): Normal distribution mean.
variance (SIMD[float64, 1]): Normal distribution variance.

Returns:

A Tensor filled with random dtype samples from Normal(mean, variance).

`ireshape`

ireshape(inout self: Self, new_shape: TensorShape)

(Inplace) Reshapes the tensor by assigning it a new shape.

Args:

new_shape (TensorShape): The new shape.

`reshape`

reshape(inout self: Self, new_shape: TensorShape) -> Self

Returns a reshaped tensor.

Args:

new_shape (TensorShape): The new shape.

Returns:

A Tensor that is a reshaped version of the original tensor.

`astype`

astype[new_type: DType](self: Self) -> Tensor[$0]

Copy the Tensor with elements cast to the new type.

Parameters:

new_type (DType): The type to cast the values to.

Returns:

A new tensor with the same values but the new type.

`clip`

clip(self: Self, lower_bound: SIMD[type, 1], upper_bound: SIMD[type, 1]) -> Self

Clips the values of the tensor.

Args:

lower_bound (SIMD[type, 1]): The lower bound.
upper_bound (SIMD[type, 1]): The upper bound.

Returns:

A clipped version of the tensor.

`unsafe_ptr`

unsafe_ptr(self: Self) -> UnsafePointer[SIMD[type, 1], 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]

Gets the underlying Data pointer to the Tensor.

Returns:

The underlying data pointer of the tensor.

`unsafe_uint8_ptr`

unsafe_uint8_ptr(self: Self) -> UnsafePointer[SIMD[uint8, 1], 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 1]

Gets the underlying Data pointer to the Tensor.

Returns:

The underlying data pointer of the tensor.

`rank`

rank(self: Self) -> Int

Gets the rank of the tensor.

Returns:

The rank of the tensor.

`num_elements`

num_elements(self: Self) -> Int

Gets the total number of elements in the tensor.

Returns:

The total number of elements in the tensor.

`bytecount`

bytecount(self: Self) -> Int

Gets the total bytecount of the tensor.

Returns:

The total bytecount of the tensor.

`spec`

spec(self: Self) -> TensorSpec

Gets the specification of the tensor.

Returns:

The underlying tensor spec of the tensor.

`shape`

shape(self: Self) -> TensorShape

Gets the shape of the tensor.

Returns:

The underlying tensor shape of the tensor.

`dim`

dim(self: Self, idx: Int) -> Int

Gets the dimension at the specified index.

Args:

idx (Int): The dimension index.

Returns:

The dimension at the specified index.

`str`

__str__(self: Self) -> String

Gets the tensor as a string.

Returns:

A compact string of the tensor.

`format_to`

format_to(self: Self, inout writer: Formatter)

Formats this Tensor to the provided formatter.

Args:

writer (Formatter): The formatter to write to.

`repr`

__repr__(self: Self) -> String

Gets the tensor as a string.

Returns:

A compact string representation of the tensor.

`load`

load[*, width: Int = 1](self: Self, index: Int) -> SIMD[type, $0]

Gets the SIMD value at the specified index.

Parameters:

width (Int): The SIMD width of the vector.

Args:

index (Int): The index of the value to retrieve.

Returns:

The SIMD value at the specified indices.

load[*, width: Int = 1](self: Self, *indices: Int) -> SIMD[type, $0]

Gets the SIMD value at the specified indices.

Parameters:

width (Int): The SIMD width of the vector.

Args:

*indices (Int): The indices of the value to retrieve.

Returns:

The SIMD value at the specified indices.

load[*, width: Int = 1](self: Self, indices: VariadicList[Int]) -> SIMD[type, $0]

Gets the SIMD value at the specified indices.

Parameters:

width (Int): The SIMD width of the vector.

Args:

indices (VariadicList[Int]): The indices of the value to retrieve.

Returns:

The SIMD value at the specified indices.

load[len: Int, /, *, width: Int = 1](self: Self, indices: StaticIntTuple[len]) -> SIMD[type, $1]

Gets the SIMD value at the specified indices.

Parameters:

len (Int): The length of the indecies.
width (Int): The SIMD width of the vector.

Args:

indices (StaticIntTuple[len]): The indices of the value to retrieve.

Returns:

The SIMD value at the specified indices.

static load(path: Path) -> Self

Read tensor from a file. The path should be a file saved with Tensor.save method.

Args:

path (Path): Path to the output file.

Returns:

The tensor read from file.

`store`

store[*, width: Int = 1](inout self: Self, index: Int, val: SIMD[type, width])

Sets the SIMD value at the specified index.

Parameters:

width (Int): The SIMD width of the vector.

Args:

index (Int): The index of the value to set.
val (SIMD[type, width]): The SIMD value to store.

store[*, width: Int = 1](inout self: Self, indices: VariadicList[Int], val: SIMD[type, width])

Sets the SIMD value at the specified indices.

Parameters:

width (Int): The SIMD width of the vector.

Args:

indices (VariadicList[Int]): The indices of the value to set.
val (SIMD[type, width]): The SIMD value to store.

store[len: Int, /, *, width: Int = 1](inout self: Self, indices: StaticIntTuple[len], val: SIMD[type, width])

Sets the SIMD value at the specified indices.

Parameters:

len (Int): The length of the indecies.
width (Int): The SIMD width of the vector.

Args:

indices (StaticIntTuple[len]): The indices of the value to set.
val (SIMD[type, width]): The SIMD value to store.

`argmax`

argmax(self: Self, *, axis: Int = -1) -> Tensor[index]

Finds the indices of the maximum element along the specified axis.

Args:

axis (Int): The axis.

Returns:

A new tensor containing the indices of the maximum elements along axis.

Raises:

Error if this Tensor's rank is larger than the maximum supported rank for argmax.

`argmin`

argmin(self: Self, *, axis: Int = -1) -> Tensor[index]

Finds the indices of the minimum element along the specified axis.

Args:

axis (Int): The axis.

Returns:

A new tensor containing the indices of the minimum elements along axis.

`tofile`

tofile(self: Self, path: Path)

Write values to a file.

Args:

path (Path): Path to the output file.

`fromfile`

static fromfile(path: Path) -> Self

Read tensor from a file.

Args:

path (Path): Path to the output file.

Returns:

The tensor read from file.

`save`

save(self: Self, path: Path)

Save given tensor to file. This method preserves shape and datatype information.

Args:

path (Path): Path of file.

Was this page helpful?

Thank you! We'll create more content like this.

Thank you for helping us improve!

😔 What went wrong?

Parameters​

Implemented traits​

Methods​

__init__​

__copyinit__​

__moveinit__​

__del__​

__getitem__​

__setitem__​

__eq__​

__ne__​

__add__​

__sub__​

__mul__​

__truediv__​

__pow__​

__radd__​

__rsub__​

__rmul__​

__rtruediv__​

__ipow__​

rand​

randn​

ireshape​

reshape​

astype​

clip​

unsafe_ptr​

unsafe_uint8_ptr​

rank​

num_elements​

bytecount​

spec​

shape​

dim​

__str__​

format_to​

__repr__​

load​

store​

argmax​

argmin​

tofile​

fromfile​

save​

Parameters

Implemented traits

Methods

`init`

`copyinit`

`moveinit`

`del`

`getitem`

`setitem`

`eq`

`ne`

`add`

`sub`

`mul`

`truediv`

`pow`

`radd`

`rsub`

`rmul`

`rtruediv`

`ipow`

`rand`

`randn`

`ireshape`

`reshape`

`astype`

`clip`

`unsafe_ptr`

`unsafe_uint8_ptr`

`rank`

`num_elements`

`bytecount`

`spec`

`shape`

`dim`

`str`

`format_to`

`repr`

`load`

`store`

`argmax`

`argmin`

`tofile`

`fromfile`

`save`