Return the number of entries in the cache.

Return the total number of samples in the dataset.

Associate the given value with the given key.

Apply a block-diagonal function: split features by sizes, transform each, concat.

Apply complex RoPE (compiled: fuses fp64 cast + view_as_complex + multiply chain).

Apply a per-token 4x4 projection matrix to feature channels grouped by 4.

Apply bidirectional (non-causal) ShortConvolution along T. Uses the forward+backward causal trick: run the causal conv in both direct

Chunk-causal ShortConvolution: global forward + per-chunk backward. Mirrors the ChunkCausalGDN recurrence semantics: * **Forward (ca

One program per (b, n, h) — processes a single (Q, K, V) head slice. Loads the first D_HALF dims as a (N_GROUPS, 4) tile (for the UCPE block-

Detach gradients from all KV cache tensors (external and module-internal). This prevents autograd from tracking historical caches across chunk

Convert x0 prediction to flow matching's prediction. x0_pred: the x0 prediction with shape [B, C, H, W] xt: the input noisy d

Efficient dequantization for the first step

Extract values from a 1-D numpy array for a batch of indices. :param arr: the 1-D numpy array. :param timesteps: a tensor of

Chunk-wise block-causal mask for video generation. Full attention within each chunk (all spatial tokens across all frames in the chunk attend

Compute tile start positions ensuring full coverage with given overlap. Returns a list of start positions such that every pixel in [0, total_