Type Alias RearrangeOp<Input, Pattern, Axes>

RearrangeOp: ResolveEinopsShape<Pattern, Input["__shape"], Axes> extends infer OutputShape
    ? OutputShape extends Shape
        ? StorageTransformation<
            "rearrange",
            TensorStorage<
                Input["__dtype"],
                OutputShape,
                ComputeStrides<OutputShape>,
                RearrangeLayout<Input["__layout"]>,
            >,
            readonly [Input],
        >
        : never & {
            __error: "Failed to resolve einops pattern";
            __inputShape: ShapeToString<Input["__shape"]>;
            __pattern: Pattern;
        }
    : never

type A = TensorStorage<Float32, [2, 3]>;
type B = RearrangeOp<A, 'h w -> w h'>; // Shape: [3, 2]
type C = RearrangeOp<A, 'h w -> (h w)'>; // Shape: [6]
Copy