torch.ReduceShape
export type ReduceShape<Pattern extends string, InputShape extends Shape> =
IsNever<SplitArrow<Pattern>> extends true
? einops_error_invalid_pattern<Pattern>
: SplitArrow<Pattern> extends {
left: infer Left extends string;
right: infer Right extends string;
}
? ParsePattern<Left> extends infer InputDims extends ParsedDim[]
? ParsePattern<Right> extends infer OutputDims extends ParsedDim[]
? HasEllipsis<InputDims> extends true
? ReduceShapeWithEllipsis<InputDims, OutputDims, InputShape>
: ReduceShapeNoEllipsis<InputDims, OutputDims, InputShape>
: DynamicShape
: DynamicShape
: einops_error_invalid_pattern<Pattern>;Patternextends stringInputShapeextends ShapeCompute the output shape of a reduce operation.
Examples
type S1 = ReduceShape<'b c h w -> b c', [2, 3, 4, 5]>;
// S1 = [2, 3]
type S2 = ReduceShape<'b c h w -> b', [2, 3, 4, 5]>;
// S2 = [2]