torch.DiagShape

export type DiagShape<S extends Shape> = S extends readonly [infer N extends number]
  ? readonly [N, N] // 1D → 2D
  : S extends readonly [infer R extends number, infer C extends number]
    ? readonly [MinDim<R, C>] // 2D → 1D (approximation)
    : DynamicShape;

Sextends Shape

Compute output shape of diag operation.

1D [n] → 2D [n, n] (creates square matrix with vector on diagonal)
2D [r, c] → 1D [min(r, c)] (extracts diagonal as vector)

Note: The actual output size depends on the diagonal offset parameter, but for type purposes we use the base case (diagonal=0).