gensbi.models.tarflow.pe#

Rotary positional embedding for the transformer flow.

Adapted from apple/ml-starflow (STARFlow); see models/tarflow/LICENSE.starflow. Faithful JAX port of the reference misc/pe.py, restricted to the paths GenSBI uses. Deliberately omitted relative to the reference: the is_1d branch (used only for the pretrained-LM top block), freqs_for='pixel'/'constant' and custom_freqs, the latent_len head-dim split (STARFlow’s text-prefix “3D” axis — GenSBI positions prefix tokens at the identity rotation instead; see docs/superpowers/specs/2026-07-11-tarflow-rope-kvcache-design.md), video duplicate handling, and the deprecated checkpoint-compat buffers.

Classes#

VisionRotaryEmbedding

2D vision RoPE with the 'lang' frequency schedule.

Functions#

apply_rope(t, freqs)

Apply a rotary embedding: t*cos(freqs) + rotate_half(t)*sin(freqs).

get_positions(h, w[, pt_seq_len])

2D patch positions in raster order (reference '2d' mode).

rotate_half(x)

Rotate adjacent channel pairs: (x1, x2) -> (-x2, x1).

Module Contents#

class gensbi.models.tarflow.pe.VisionRotaryEmbedding(dim, pt_seq_len=16, theta=10000)[source]#

Bases: flax.nnx.Module

2D vision RoPE with the 'lang' frequency schedule.

Port of the reference VisionRotaryEmbeddingFast without the latent_len split: each of the two position axes is rotated with the same 1 / theta**(2i/dim) frequency table, and the concatenated per-axis angles are repeated pairwise to cover 2*dim channels (= the full head dimension when dim = head_dim // 2).

Parameters:

  • dim (int) – Half the attention head dimension. Must be even.

  • pt_seq_len (int, optional) – Pre-training sequence length recorded for position building. Default is 16.

  • theta (int, optional) – Frequency base. Default is 10000.

__call__(pos)[source]#

Build rotation angles for 2D positions.

Parameters:

pos (Array) – Positions of shape (..., 2).

Returns:

Angles of shape (..., 2*dim).

Return type:

Array

freqs[source]#
pt_seq_len = 16[source]#
gensbi.models.tarflow.pe.apply_rope(t, freqs)[source]#

Apply a rotary embedding: t*cos(freqs) + rotate_half(t)*sin(freqs).

Parameters:

  • t (Array) – Tensor to rotate; last dimension even.

  • freqs (Array) – Rotation angles, broadcastable to t’s shape.

Returns:

Rotated tensor, same shape as t.

Return type:

Array

gensbi.models.tarflow.pe.get_positions(h, w, pt_seq_len=None)[source]#

2D patch positions in raster order (reference '2d' mode).

Coordinates are normalized by sqrt(h*w) and rescaled to pt_seq_len (the reference’s resolution-transfer schedule). For a square grid with pt_seq_len == h == w this reduces to plain integer coordinates.

Parameters:

  • h (int) – Patch-grid height and width.

  • w (int) – Patch-grid height and width.

  • pt_seq_len (int or None, optional) – Pre-training sequence length; defaults to sqrt(h*w).

Returns:

Positions of shape (h*w, 2).

Return type:

Array

gensbi.models.tarflow.pe.rotate_half(x)[source]#

Rotate adjacent channel pairs: (x1, x2) -> (-x2, x1).

Port of the reference rotate_half (einops (d r) with r=2 means adjacent pairs).

Parameters:

x (Array) – Input whose last dimension is even.

Returns:

Same shape as x.

Return type:

Array