torch.SoftmaxShape

export type SoftmaxShape<S extends Shape, Dim extends number> = number extends Dim
  ? S
  : number extends S['length']
    ? S
    : LessThan<Dim, S['length']> extends true
      ? S
      : softmax_error_dim_out_of_range<Dim, S['length']>;

Sextends Shape

Dimextends number

Compute output shape of softmax with bounds checking. softmax preserves shape but requires dim < rank.