Mojo module

managed_tensor_slice

ManagedTensorSlice and the types kernel authors compose against.

A custom kernel's entry-point signature uses these:

ManagedTensorSlice — view of a tensor argument.
IOSpec (and IO) — input/output/mutability annotations.
RuntimeTensorSpec / StaticTensorSpec — runtime + compile-time tensor metadata.
Fusion traits (InputFusion, OutputFusion, ...) and their _NoFusion* sentinels.

The decorators that register a kernel (register, register_internal, view_kernel) live next to this file in register.mojo.

`comptime` values

`DynamicTensor`

comptime DynamicTensor[dtype: DType, rank: Int] = ManagedTensorSlice[IOUnknown, static_spec=_get_unknown_tensor_spec[dtype, rank]()]

Parameters

dtype (DType):
rank (Int):

`FusedInput`

comptime FusedInput = IOSpec()

`FusedOutput`

comptime FusedOutput = IOSpec()

`Input`

comptime Input = IOSpec()

`InputTensor`

comptime InputTensor = ManagedTensorSlice[Input, static_spec=?]

`InputVariadicTensors`

comptime InputVariadicTensors = VariadicTensors[?, Input, static_specs=?]

`IOUnknown`

comptime IOUnknown = IOSpec()

`MutableInput`

comptime MutableInput = IOSpec()

`Output`

comptime Output = IOSpec()

`OutputTensor`

comptime OutputTensor = ManagedTensorSlice[Output, static_spec=?]

`OutputVariadicTensors`

comptime OutputVariadicTensors = VariadicTensors[?, Output, static_specs=?]

Structs

IO:
IOSpec: Parameter used to encode whether a particular tensor argument to a DPS kernel is an output, input, or mutable input.
ManagedTensorSlice: A view of a tensor that does not own the underlying allocated pointer. When the object lifetime ends it does not free the underlying pointer. Conversely, if a ManagedTensorSlice is created, it will not extend the life of the underlying pointer.
RuntimeTensorSpec:
StaticTensorSpec:
StaticTensorSpecInternal:
StaticTensorSpecList: A statically indexable list of data that can be assembled into a StaticTensorSpecList on demand. This handles the complexities that arise with heterogenous specs.
VariadicTensors: A tuple-like container of tensors representing variadic arguments from the graph compiler.

Traits

ComputeOutputFusion: Trait for compute-output fusion structs that transform values before storing.
ElementwiseFusion: Trait for pure elementwise fusion structs emitted by the graph compiler.
InputFusion: Trait for input fusion structs that provide custom load behavior.
OutputFusion: Trait for output fusion structs that provide custom store behavior.

Functions

foreach: Apply the function func to each element of the tensor slice.
get_row_major_tensor_spec_static: Returns a row-major StaticTensorSpec from compile-time Int dimensions.
trace_slice_arg: Helper to stringify the type and shape of a kernel argument for tracing.

comptime values​

DynamicTensor​

Parameters​

FusedInput​

FusedOutput​

Input​

InputTensor​

InputVariadicTensors​

IOUnknown​

MutableInput​

Output​

OutputTensor​

OutputVariadicTensors​

Structs​

Traits​

Functions​