For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

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.
• 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[IOSpec[_, _].Unknown, static_spec=_get_unknown_tensor_spec[dtype, rank]()]

Parameters

• ​dtype (DType):
• ​rank (Int):

InputTensor

comptime InputTensor = ManagedTensorSlice[IOSpec[_, _].Input, static_spec=_]

InputVariadicTensors

comptime InputVariadicTensors = VariadicTensors[_, IOSpec[_, _].Input, static_specs=_]

OutputTensor

comptime OutputTensor = ManagedTensorSlice[IOSpec[_, _].Output, static_spec=_]

OutputVariadicTensors

comptime OutputVariadicTensors = VariadicTensors[_, IOSpec[_, _].Output, static_specs=_]

Structs

• ​IO: Tags the direction and fusion kind of a tensor argument to a DPS kernel.
• ​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.
• ​StaticTensorSpec: Carries the compile-time and runtime metadata describing a tensor argument.
• ​StaticTensorSpecInternal: Stores the runtime-only portion of a StaticTensorSpec.
• ​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.
• ​ComputeOutputFusionTile: Trait for tile-based compute-output fusion structs that transform a TileTensor before storing.
• ​ElementwiseFusion: Trait for pure elementwise fusion structs emitted by the graph compiler.
• ​ElementwiseFusionTile: The TileTensor variant of ElementwiseFusion: a tile-based pure elementwise fusion struct 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.
• ​OutputFusionTile: Trait for tile-based output fusion structs that provide custom store behavior for a TileTensor.

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.