vortex_torch.cache.reshape

Classes

Reshape(batch_dim, x2, y2)

Same-numel reshape of a cache block's inner two axes.
class Reshape(batch_dim, x2, y2)[source]

Bases: vOp

Same-numel reshape of a cache block’s inner two axes.

Math:
\[X\in\mathbb{R}^{1\times x_1\times y_1} \;\longrightarrow\; Y\in\mathbb{R}^{1\times x_2\times y_2},\qquad x_2\,y_2 = x_1\,y_1,\]

reading the flat \(x_1 y_1\) elements row-major into the new \((x_2, y_2)\) layout (a Triton-tile tl.reshape(); no data movement beyond the existing load/store).

__init__:

Reshape(-1, x2, y2) — the leading dim must be -1 (the per-block batch axis, always 1); x2*y2 must equal the input’s x1*y1 (checked at trace time).

__call__:

y = op(x, output, loc=loc, ctx=ctx) — x [1, x_1, y_1] → [1, x_2, y_2]. PAGED iff a PAGED output is supplied, else RAGGED.

Parameters: