Skip to main content
Log in

Mojo struct

Tensor

An owned, indexible buffer type.

Fields

  • name (Optional[String]):

Implemented traits

AnyType, CollectionElement, Copyable, Formattable, Movable, TensorLike

Methods

__init__

__init__(inout self: Self)

Default constructor for Tensor. Accessing the elements of default constructed tensor is undefined behavior.

__init__(inout self: Self, *, owned device_tensor: DeviceTensor)

Creates a tensor from DeviceTensor.

Args:

  • device_tensor (DeviceTensor): DeviceTensor to create tensor from.

__init__(inout self: Self, shape: TensorShape, device: Optional[Device] = #kgen.none)

Creates tensor with given shape on the given device. If device is not given tensor will be created on cpu.

Args:

  • shape (TensorShape): Shape of the tensor.
  • device (Optional[Device]): Device on which tensor is to be allocated.

__init__(inout self: Self, tensor: Tensor[type])

Converts max.tensor to max.driver.Tensor. This creates tensor on the CPU.

Args:

  • tensor (Tensor[type]): Tensor to copy from.

__moveinit__

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

Move constructor for Tensor.

Args:

  • existing (Self): Instance to move from.

__del__

__del__(owned self: Self)

Destructor for the tensor.

__getitem__

__getitem__(inout self: Self, *indices: Int) -> ref [*[0,0]] 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__(ref [self_is_lifetime] self: Self, *slices: Slice) -> TensorSlice[$0, type, rank, $1]

Returns a view of the tensor conforming to given slices. If given a single slice : the view would point to the entire tensor. The returned slice has the lifetime of tensor and will extend the lifetime of tensor accordingly.

Args:

  • *slices (Slice): Dimension slices to slice against.

Returns:

View of the tensor according to given slices.

spec

spec(self: Self) -> TensorSpec

Gets the spec of tensor.

Returns Spec of the tensor.

unsafe_slice

unsafe_slice(self: Self, *slices: Slice) -> ManagedTensorSlice[type, rank]

Returns a view of the tensor conforming to given slices. If given a single slice : the view would point to the entire tensor. The caller is responsible to make sure tensor outlives the returned slice.

Args:

  • *slices (Slice): Dimension slices to slice against.

Returns:

View of the tensor according to given slices.

to_device_tensor

to_device_tensor(owned self: Self) -> DeviceTensor

Converts the tensor to a DeviceTensor.

Returns:

DeviceTensor pointing to the memory owned by tensor.

unsafe_ptr

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

Gets a pointer to the underlying memory.

Note: The caller is responsible for ensuring that the returned pointer is not used after it's owner is last used.

Parameters:

  • __type (DType): If given the pointer will be rebound to this type. Defaulted to type of tensor.

Returns:

Pointer to the beginning of tensor data.

take

take(inout self: Self) -> Self

Takes self's resources and replaces them with default initialized values.

Returns:

An instance of tensor.

__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.

move_to

move_to(owned self: Self, device: Device) -> Self

Returns self if already allocated on device, otherwise copy the contents of self to device.

Args:

  • device (Device): The Device of the returned buffer.

Returns:

Instance of Tensor allocated on given device.

Was this page helpful?