torch.RepeatShape
export type RepeatShape<
Pattern extends string,
InputShape extends Shape,
Axes extends AxesRecord = Record<string, never>,
> =
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
? RepeatShapeWithEllipsis<InputDims, OutputDims, InputShape, Axes>
: RepeatShapeNoEllipsis<InputDims, OutputDims, InputShape, Axes>
: DynamicShape
: DynamicShape
: einops_error_invalid_pattern<Pattern>;Patternextends stringInputShapeextends ShapeAxesextends AxesRecordCompute the output shape of a repeat operation.
Examples
type S1 = RepeatShape<'h w -> h w c', [4, 5], { c: 3 }>;
// S1 = [4, 5, 3]
type S2 = RepeatShape<'h w -> (h 2) (w 3)', [4, 5], {}>;
// S2 = [8, 15]