pub struct Shape { /* private fields */ }Expand description
Tensor shape: ordered list of dimensions + element type.
SmallVec<[Dim; 4]> avoids heap allocation for up to 4D tensors (the common case).
Implementations§
Source§impl Shape
impl Shape
Source§impl Shape
impl Shape
Sourcepub fn from_dims(dims: &[Dim], dtype: DType) -> Self
pub fn from_dims(dims: &[Dim], dtype: DType) -> Self
Create a shape with mixed static/dynamic dimensions.
pub fn rank(&self) -> usize
pub fn dtype(&self) -> DType
pub fn dims(&self) -> &[Dim]
pub fn dim(&self, i: usize) -> Dim
Sourcepub fn gather_index_count(&self, axis: usize) -> usize
pub fn gather_index_count(&self, axis: usize) -> usize
How many positions a gather/scatter along axis takes, when self is
the shape of the INDEX operand.
Gather’s output rank is data_rank - 1 + index_rank, so an index may be
either:
- rank >
axis— the usual case, where the index carries a dimension at the gather axis and that dimension is the count; or - a flat list (typically rank 1) that has no dimension at
axisand whose whole extent IS the count.
Every backend previously wrote idx_shape.dim(axis) here, which panics
with “Shape::dim(1) out of bounds for rank 1” on the second form — so a
rank-1 gather built a correct forward graph and then blew up the moment
it needed a backward pass. Sharing the rule keeps the five backends that
need it from disagreeing about it.
Sourcepub fn dynamic_symbols(&self) -> Vec<u32>
pub fn dynamic_symbols(&self) -> Vec<u32>
Set of dynamic dim symbols this shape references. Useful for “what bindings does this graph need?” queries on inputs.
Sourcepub fn bind(&self, bindings: &DimBinding) -> Self
pub fn bind(&self, bindings: &DimBinding) -> Self
Specialize the shape against a binding (symbol → static size). Unknown symbols stay Dim::Dynamic. Plan #54: the
step that takes a “compile once, run at any seq length” graph
and produces the runtime-specific concrete shape.
Sourcepub fn num_elements(&self) -> Option<usize>
pub fn num_elements(&self) -> Option<usize>
Total number of elements (only if all dims are static).
Sourcepub fn size_bytes(&self) -> Option<usize>
pub fn size_bytes(&self) -> Option<usize>
Total size in bytes (only if all dims are static).
Sourcepub fn with_dtype(self, dtype: DType) -> Self
pub fn with_dtype(self, dtype: DType) -> Self
Change dtype (for cast operations).