IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

tensor_arg_traits

Composable traits for graph-compiler kernel tensor arguments.

A kernel declares each tensor argument by combining small, orthogonal traits with trait composition instead of naming a single monolithic tensor type.

The traits fall into three groups:

  • Role traits (Input, Output, MutableInput) declare how the kernel uses the argument: read-only, write-only destination, or read-write in place. Exactly one role applies to each argument.
  • The Fused marker declares that the operand participates in prologue or epilogue fusion.
  • Metadata and access traits describe what the kernel can read off the argument (Tensor, DenseTensor) and how it touches individual elements.

Element access is split into marker / capability pairs. A graph-visible marker (LoadAccess, StoreAccess, TransformAccess) records the access form in the distilled contract but adds no method; the matching capability (Loadable, Storable, Transformable) additionally supplies the method the kernel body calls. A kernel typically binds the marker in its signature and refines to the capability inside its body, so the same argument type works whether it stays on the host or is projected to a device.

Traits

  • DenseTensor: Metadata trait: a Tensor backed by a dense (contiguous) layout.
  • Fused: Marker trait: the operand participates in fusion.
  • Input: Tells the graph compiler this argument is a read-only input.
  • Loadable: LoadAccess refined with the callable load method.
  • LoadAccess: Marker trait: this tensor supports SIMD loads.
  • MutableInput: Tells the graph compiler this argument is a mutable input.
  • Output: Tells the graph compiler this argument is an output.
  • Storable: StoreAccess refined with the callable store method.
  • StoreAccess: Marker trait: this tensor supports SIMD stores.
  • Tensor: Metadata trait exposing a tensor's element type, known rank, and shape.
  • TileTensorable: A DenseTensor that can project to a TileTensor view over its storage.
  • Transformable: TransformAccess refined with the callable transform method.
  • TransformAccess: Marker trait: this tensor transforms a kernel-computed SIMD value.