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 struct

StaticTensorSpec

struct StaticTensorSpec[dtype: DType, rank: Int, static_layout: TensorLayout, InFusion: InputFusion = _NoFusionIn, OutFusion: OutputFusion = _NoFusionOut, ComputeFusion: ComputeOutputFusion = _NoComputeFusion, ComputeFusionTile: ComputeOutputFusionTile = _NoComputeFusionTile, OutFusionTile: OutputFusionTile = _NoOutputFusionTile]

Carries the compile-time and runtime metadata describing a tensor argument.

The compile-time parameters encode the element dtype, tensor rank, static layout, and optional fusion traits. The runtime fields store the alignment and address space of the backing memory. Custom kernels receive ManagedTensorSlice instances parameterized by a StaticTensorSpec.

Fields

  • alignment (Int):
  • address_space (AddressSpace):

Implemented traits

AnyType, Copyable, Deinitable, ImplicitlyCopyable, Movable

comptime members

get_unknown

comptime get_unknown = fn_literal

shape_tuple

comptime shape_tuple = coord_to_int_tuple[static_layout._shape_types]()

static_size

comptime static_size = Layout(coord_to_int_tuple[static_layout._shape_types](), coord_to_int_tuple[static_layout._stride_types]()).size()

strides_tuple

comptime strides_tuple = coord_to_int_tuple[static_layout._stride_types]()

Methods

__init__

def __init__(out self, alignment: Int, address_space: AddressSpace)

def __init__(out self, internals: StaticTensorSpecInternal[dtype, rank])

Returns a StaticTensorSpec from a StaticTensorSpecInternal.

to_unfused

def to_unfused(self) -> StaticTensorSpec[dtype, rank, static_layout]

Returns a copy with sentinel (no-op) fusion types.

The runtime fields (alignment, etc.) are identical; only the compile-time fusion type parameters change.

Returns:

StaticTensorSpec[dtype, rank, static_layout]

with_tile_layout

def with_tile_layout[new_layout: TensorLayout](self) -> StaticTensorSpec[dtype, new_layout.rank, new_layout]

Returns:

StaticTensorSpec[dtype, new_layout.rank, new_layout]

def with_tile_layout[new_rank: Int, new_layout: TensorLayout](self) -> StaticTensorSpec[dtype, new_rank, new_layout]

Returns:

StaticTensorSpec[dtype, new_rank, new_layout]

with_tile_layout_and_alignment

def with_tile_layout_and_alignment[new_layout: TensorLayout](self, new_alignment: Int) -> StaticTensorSpec[dtype, new_layout.rank, new_layout]

Returns:

StaticTensorSpec[dtype, new_layout.rank, new_layout]

def with_tile_layout_and_alignment[new_rank: Int, new_layout: TensorLayout](self, new_alignment: Int) -> StaticTensorSpec[dtype, new_rank, new_layout]

Returns:

StaticTensorSpec[dtype, new_rank, new_layout]

with_int_tuple_layout

def with_int_tuple_layout[new_rank: Int, new_shape: IntTuple, new_strides: IndexList[new_strides.size, element_type=new_strides.element_type]](self) -> StaticTensorSpec[dtype, new_rank, Layout[*?, *?]]

Returns:

StaticTensorSpec[dtype, new_rank, Layout[*?, *?]]

with_int_tuple_layout_and_alignment

def with_int_tuple_layout_and_alignment[new_rank: Int, new_shape: IntTuple, new_strides: IndexList[new_strides.size, element_type=new_strides.element_type]](self, new_alignment: Int) -> StaticTensorSpec[dtype, new_rank, Layout[*?, *?]]

Returns:

StaticTensorSpec[dtype, new_rank, Layout[*?, *?]]

with_row_major_int_tuple_layout

def with_row_major_int_tuple_layout[new_rank: Int, new_shape: IntTuple](self) -> StaticTensorSpec[dtype, new_rank, Layout[*?, *?]]

Returns:

StaticTensorSpec[dtype, new_rank, Layout[*?, *?]]

with_input_fusion

def with_input_fusion[F: InputFusion](self) -> StaticTensorSpec[dtype, rank, static_layout, F, OutFusion, ComputeFusion, ComputeFusionTile, OutFusionTile]

Returns:

StaticTensorSpec[dtype, rank, static_layout, F, OutFusion, ComputeFusion, ComputeFusionTile, OutFusionTile]

with_output_fusion

def with_output_fusion[F: OutputFusion](self) -> StaticTensorSpec[dtype, rank, static_layout, InFusion, F, ComputeFusion, ComputeFusionTile, OutFusionTile]

Returns:

StaticTensorSpec[dtype, rank, static_layout, InFusion, F, ComputeFusion, ComputeFusionTile, OutFusionTile]

with_compute_fusion

def with_compute_fusion[F: ComputeOutputFusion](self) -> StaticTensorSpec[dtype, rank, static_layout, InFusion, OutFusion, F, ComputeFusionTile, OutFusionTile]

Returns:

StaticTensorSpec[dtype, rank, static_layout, InFusion, OutFusion, F, ComputeFusionTile, OutFusionTile]

with_compute_fusion_tile

def with_compute_fusion_tile[F: ComputeOutputFusionTile](self) -> StaticTensorSpec[dtype, rank, static_layout, InFusion, OutFusion, ComputeFusion, F, OutFusionTile]

Returns:

StaticTensorSpec[dtype, rank, static_layout, InFusion, OutFusion, ComputeFusion, F, OutFusionTile]

with_output_fusion_tile

def with_output_fusion_tile[F: OutputFusionTile](self) -> StaticTensorSpec[dtype, rank, static_layout, InFusion, OutFusion, ComputeFusion, ComputeFusionTile, F]

Returns:

StaticTensorSpec[dtype, rank, static_layout, InFusion, OutFusion, ComputeFusion, ComputeFusionTile, F]

to_layout

def to_layout(self) -> Layout

Returns:

Layout

get_internals

def get_internals(self) -> StaticTensorSpecInternal[dtype, rank]

Returns a StaticTensorSpecInternal from a StaticTensorSpec.

Returns:

StaticTensorSpecInternal[dtype, rank]