torch.AtShape
export type AtShape<
S extends Shape,
Specs extends readonly IndexSpec[],
> = number extends S['length'] ? DynamicShape : ApplyIndexSpecs<S, Specs>;Sextends ShapeSpecsextends readonly IndexSpec[]Compute output shape of .at() multi-dimensional indexing.
Examples
// Basic indexing
// tensor.at(0) on [2, 3, 4] → [3, 4]
type R1 = AtShape<readonly [2, 3, 4], readonly [0]>;
// Range slicing
// tensor.at(null, [1, 3]) on [2, 5, 4] → [2, 2, 4]
type R2 = AtShape<readonly [2, 5, 4], readonly [null, readonly [1, 3]]>;
// Mixed indexing
// tensor.at(0, [1, 3], null) on [2, 5, 4] → [2, 4]
type R3 = AtShape<readonly [2, 5, 4], readonly [0, readonly [1, 3], null]>;
// Ellipsis - select last dimension
// tensor.at('...', 0) on [2, 3, 4] → [2, 3]
type R4 = AtShape<readonly [2, 3, 4], readonly ['...', 0]>;
// Ellipsis - select first and last
// tensor.at(0, '...', 0) on [2, 3, 4, 5] → [3, 4]
type R5 = AtShape<readonly [2, 3, 4, 5], readonly [0, '...', 0]>;