torch.SelectShape

export type SelectShape<
  S extends Shape,
  Dim extends number,
  Index extends number = 0,
> = number extends S['length']
  ? DynamicShape
  : number extends Dim
    ? DynamicShape
    : number extends Index
      ? DynamicShape
      : GetDimSize<S, Dim> extends infer DimSize extends number
        ? LessThan<Index, DimSize> extends true
          ? RemoveDim<S, Dim> // Valid! Return shape with dimension removed
          : select_error_index_out_of_bounds<Index, DimSize, Dim>
        : DynamicShape;

Sextends Shape

Dimextends number

Indexextends number

Compute output shape of select operation with compile-time bounds checking. select(dim, index) removes dimension dim from the shape.

Validates:

index >= 0 (implicit, as index is a positive literal)
index < dimSize

Examples

// Valid select
type R1 = SelectShape<readonly [2, 3, 4], 1, 1>; // readonly [2, 4]

// Error: index out of bounds
type R2 = SelectShape<readonly [2, 3, 4], 1, 5>;
// select_error_index_out_of_bounds<5, 3, 1>