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
StableTensor
struct StableTensor[mut: Bool, input: IO, dtype: DType, rank: Int, //, io_spec: IOSpec[mut, input], *, static_spec: StaticTensorSpec[dtype, rank, static_spec.static_layout]]
A tensor view whose data pointer lives in a stable address slot.
The view holds the slot's address and the tensor's layout; the data pointer
is read out of the slot on each access. A device graph that records this
view bakes in the slot, so
DeviceGraphBuilder.place_stable()
can point the slot at a different tensor before each replay without the
graph being rebuilt.
The slot and the pointer it holds are both device addresses, so the view
is built on the host but read only on the device: pass it to a kernel and
convert to a ManagedTensorSlice or TileTensor there. The conversions
are compile errors on the host.
Parameters
- mut (
Bool): Whether the tensor is mutable. - input (
IO): The IO kind of the tensor. - dtype (
DType): The element type of the tensor. - rank (
Int): The rank of the tensor. - io_spec (
IOSpec[mut, input]): The IO specification of the tensor. - static_spec (
StaticTensorSpec[dtype, rank, static_spec.static_layout]): The static specification of the tensor, including its layout.
Implemented traits
AnyType,
Copyable,
Deinitable,
DevicePassable,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
device_type
comptime device_type = StableTensor[io_spec, static_spec=static_spec]
RuntimeLayout
comptime RuntimeLayout = Layout[static_layout._shape_types, static_layout._stride_types]
Methods
get_type_name
shape
def shape(self) -> IndexList[rank]
Gets the shape of this tensor, as an IndexList.
Host-callable: the shape lives in the view, not behind the slot.
Returns:
IndexList[rank]: The shape of this tensor.
size
def size(self) -> Int
Computes the tensor's number of elements.
Host-callable: the size lives in the view, not behind the slot.
Returns:
Int: The total number of elements in the tensor.
unsafe_ptr
def unsafe_ptr(self) -> Pointer[Scalar[dtype], MutUntrackedOrigin]
Reads the data pointer the slot holds at the time of the call.
Device only: the slot is device memory.
Returns:
Pointer[Scalar[dtype], MutUntrackedOrigin]: The tensor's current data pointer.
to_tile_tensor
def to_tile_tensor(self) -> TileTensor[dtype, Layout[static_spec.static_layout._shape_types, static_spec.static_layout._stride_types], MutUntrackedOrigin, Engine=StableAddrEngine]
Builds a TileTensor that keeps reading through the stable address.
The result stores the stable address rather than resolving it at conversion time, so the conversion may run on the host while the slot is still dereferenced only on the device. The cost is one additional pointer dereference per load or store.
Returns:
TileTensor[dtype, Layout[static_spec.static_layout._shape_types, static_spec.static_layout._stride_types], MutUntrackedOrigin, Engine=StableAddrEngine]: A slot-backed tile tensor with this view's layout.