vortex_torch.cache.reduce_interleave¶

Classes

`L2NormInterleave`([dim, k])	Interleaved L2-norm over groups of `k` (a `ReduceInterleave`).
`MaxInterleave`([dim, k])	Interleaved max over groups of `k` (a `ReduceInterleave`).
`MeanInterleave`([dim, k])	Interleaved mean over groups of `k` (a `ReduceInterleave`).
`MinInterleave`([dim, k])	Interleaved min over groups of `k` (a `ReduceInterleave`).
`ReduceInterleave`([dim, k])	Interleaved 1-D reduction — folds consecutive groups of `k` elements.

class ReduceInterleave(dim=1, k=2)[source]¶

Bases: vOp

Interleaved 1-D reduction — folds consecutive groups of k elements.

Math:

For input \(X\in\mathbb{R}^{B\times N\times D}\) and a per-group reduction \(\rho\) (mean / max / min / L2-norm, fixed by the subclass), each group of \(k\) adjacent elements collapses to one:

\[\begin{split}\begin{aligned} (\text{dim}=1,\ N\%k=0):\quad & Y_{b,m,d} = \rho_{\,0 \le j < k}\, X_{b,\,mk+j,\,d}, \\ (\text{dim}=2,\ D\%k=0):\quad & Y_{b,n,e} = \rho_{\,0 \le j < k}\, X_{b,n,\,ek+j}. \end{aligned}\end{split}\]

__init__:

ReduceInterleave(dim=1, k=2) — inner axis (1 over \(N\), 2 over \(D\)) and group size k (must divide the reduced axis; k=1 is identity, k = full length recovers the plain reduction).

__call__:

op(x, output, loc=loc, ctx=ctx) — x [B, N, D]; the reduced axis shrinks by a factor of k (N→N/k or D→D/k). PAGED iff a PAGED output is supplied, else RAGGED.

Note:

use a concrete subclass — MeanInterleave, MaxInterleave, MinInterleave, L2NormInterleave.

Parameters:

dim (int)
k (int)

class MeanInterleave(dim=1, k=2)[source]¶

Bases: ReduceInterleave

Interleaved mean over groups of k (a ReduceInterleave).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{b,m,d} = \frac{1}{k}\sum_{j=0}^{k-1} X_{b,\,mk+j,\,d}, \\ (\text{dim}=2):\quad & Y_{b,n,e} = \frac{1}{k}\sum_{j=0}^{k-1} X_{b,n,\,ek+j}. \end{aligned}\end{split}\]

__init__:

MeanInterleave(dim=1, k=2) — inner axis (1/2) and group size k.

Parameters:

dim (int)
k (int)

class MaxInterleave(dim=1, k=2)[source]¶

Bases: ReduceInterleave

Interleaved max over groups of k (a ReduceInterleave).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{b,m,d} = \max_{0 \le j < k} X_{b,\,mk+j,\,d}, \\ (\text{dim}=2):\quad & Y_{b,n,e} = \max_{0 \le j < k} X_{b,n,\,ek+j}. \end{aligned}\end{split}\]

__init__:

MaxInterleave(dim=1, k=2) — inner axis (1/2) and group size k.

Parameters:

dim (int)
k (int)

class MinInterleave(dim=1, k=2)[source]¶

Bases: ReduceInterleave

Interleaved min over groups of k (a ReduceInterleave).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{b,m,d} = \min_{0 \le j < k} X_{b,\,mk+j,\,d}, \\ (\text{dim}=2):\quad & Y_{b,n,e} = \min_{0 \le j < k} X_{b,n,\,ek+j}. \end{aligned}\end{split}\]

__init__:

MinInterleave(dim=1, k=2) — inner axis (1/2) and group size k.

Parameters:

dim (int)
k (int)

class L2NormInterleave(dim=1, k=2)[source]¶

Bases: ReduceInterleave

Interleaved L2-norm over groups of k (a ReduceInterleave).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{b,m,d} = \Big(\sum_{j=0}^{k-1} X_{b,\,mk+j,\,d}^2\Big)^{1/2}, \\ (\text{dim}=2):\quad & Y_{b,n,e} = \Big(\sum_{j=0}^{k-1} X_{b,n,\,ek+j}^2\Big)^{1/2}. \end{aligned}\end{split}\]

__init__:

L2NormInterleave(dim=1, k=2) — inner axis (1/2) and group size k.

Note:

a pure \(L_2\) norm per group (no division) — not RMS.

Parameters:

dim (int)
k (int)