vortex_torch.indexer.reduce

Classes

L2Norm([dim])	L2-norm reduction over one logical axis (a Reduce).
Max([dim])	Max reduction over one logical axis (a Reduce).
Mean([dim])	Mean reduction over one logical axis (a Reduce).
Min([dim])	Min reduction over one logical axis (a Reduce).
Reduce([dim])	Generic 1-D reduction over one axis of a rank-3 logical tensor.
Sum([dim])	Sum reduction over one logical axis (a Reduce).

class Reduce(dim=1)

Bases: vOp

Generic 1-D reduction over one axis of a rank-3 logical tensor.

Math:

For input \(X\in\mathbb{R}^{N\times D_0\times D_1}\) and a per-axis reduction \(\rho\) (mean / max / min / L2-norm / sum, fixed by the subclass):

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{n,0,d} = \rho_{\,0 \le i < D_0}\, X_{n,i,d}, \\ (\text{dim}=2):\quad & Y_{n,d,0} = \rho_{\,0 \le j < D_1}\, X_{n,d,j}. \end{aligned}\end{split}\]

dim=0 collapses the packed leading axis to one row per (batch, kv\_head).

__init__:

Reduce(dim=1) — logical axis to reduce, one of 0 / 1 / 2.

__call__:

y = op(x, ctx=ctx) — x is [N, D_0, D_1]; the reduced axis is kept with size 1. For dim ∈ {1, 2} the output is BATCHED iff the input is; dim=0 requires RAGGED input and returns BATCHED.

Note:

Use a concrete subclass — Max, Min, Mean, L2Norm, Sum.

Parameters:

dim (int)

class Max(dim=1)

Bases: Reduce

Max reduction over one logical axis (a Reduce).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{n,0,d} = \max_{0 \le i < D_0} X_{n,i,d}, \\ (\text{dim}=2):\quad & Y_{n,d,0} = \max_{0 \le j < D_1} X_{n,d,j}. \end{aligned}\end{split}\]

__init__:

Max(dim=1) — axis to reduce (1 → \(D_0\), 2 → \(D_1\)).

__call__:

y = op(x, ctx=ctx) — [N, D_0, D_1] → [N, 1, D_1] (dim=1) or [N, D_0, 1] (dim=2).

Parameters:

dim (int)

class Min(dim=1)

Bases: Reduce

Min reduction over one logical axis (a Reduce).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{n,0,d} = \min_{0 \le i < D_0} X_{n,i,d}, \\ (\text{dim}=2):\quad & Y_{n,d,0} = \min_{0 \le j < D_1} X_{n,d,j}. \end{aligned}\end{split}\]

__init__:

Min(dim=1) — axis to reduce (1 → \(D_0\), 2 → \(D_1\)).

__call__:

y = op(x, ctx=ctx) — [N, D_0, D_1] → [N, 1, D_1] (dim=1) or [N, D_0, 1] (dim=2).

Parameters:

dim (int)

class Mean(dim=1)

Bases: Reduce

Mean reduction over one logical axis (a Reduce).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{n,0,d} = \frac{1}{D_0}\sum_{i=0}^{D_0-1} X_{n,i,d}, \\ (\text{dim}=2):\quad & Y_{n,d,0} = \frac{1}{D_1}\sum_{j=0}^{D_1-1} X_{n,d,j}. \end{aligned}\end{split}\]

__init__:

Mean(dim=1) — axis to reduce (1 → \(D_0\), 2 → \(D_1\)).

__call__:

y = op(x, ctx=ctx) — [N, D_0, D_1] → [N, 1, D_1] (dim=1) or [N, D_0, 1] (dim=2).

Parameters:

dim (int)

class L2Norm(dim=1)

Bases: Reduce

L2-norm reduction over one logical axis (a Reduce).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{n,0,d} = \Big(\sum_{i=0}^{D_0-1} X_{n,i,d}^2\Big)^{1/2}, \\ (\text{dim}=2):\quad & Y_{n,d,0} = \Big(\sum_{j=0}^{D_1-1} X_{n,d,j}^2\Big)^{1/2}. \end{aligned}\end{split}\]

__init__:

L2Norm(dim=1) — axis to reduce (1 → \(D_0\), 2 → \(D_1\)).

__call__:

y = op(x, ctx=ctx) — [N, D_0, D_1] → [N, 1, D_1] (dim=1) or [N, D_0, 1] (dim=2).

Parameters:

dim (int)

class Sum(dim=1)

Bases: Reduce

Sum reduction over one logical axis (a Reduce).

Math:

\[\begin{split}\begin{aligned} (\text{dim}=1):\quad & Y_{n,0,d} = \sum_{i=0}^{D_0-1} X_{n,i,d}, \\ (\text{dim}=2):\quad & Y_{n,d,0} = \sum_{j=0}^{D_1-1} X_{n,d,j}. \end{aligned}\end{split}\]

__init__:

Sum(dim=1) — axis to reduce (1 → \(D_0\), 2 → \(D_1\)).

__call__:

y = op(x, ctx=ctx) — [N, D_0, D_1] → [N, 1, D_1] (dim=1) or [N, D_0, 1] (dim=2).

Parameters:

dim (int)