chunking

Utilities for sequential chunked forward passes of causal noise layers.

A chunked forward pass transforms a sequence one fixed-size chunk at a time, so that a seeded noise layer can be prefix-invariant: Appending tokens never changes the transformed outputs of the tokens already there. This module holds the prefix-stable noise sampling such a pass relies on, whose values at each position depend only on the seed and the position's chunk, and the splitting of a sequence into padded chunks.

Classes:

Name	Description
PaddedChunk	One chunk of a sequence, padded at the end to the chunk size.

Functions:

Name	Description
derive_chunk_seed	Derive the seed of the generator that samples the noise of one chunk of a chunked forward pass.
padded_chunks	Split tensors along their sequence dimension into consecutive chunks, each padded at the end to chunk_size positions.
sample_chunked_noise	Sample standard normal noise whose values at each position depend only on seed and the position's chunk.

PaddedChunk dataclass

One chunk of a sequence, padded at the end to the chunk size.

Added in version v4.3.0.

Methods:

Name	Description
unpad	Drop the padded positions from a tensor computed for this chunk.

Attributes:

Name	Type	Description
index	tuple[slice, ...]	The index of the chunk's real positions in a tensor of the whole sequence, e.g. output[chunk.index] = ....
positions	slice	The slice of the sequence's real positions that the chunk covers. It is shorter than the chunk size only for a padded last chunk.
tensors	tuple[Tensor, ...]	The padded chunk of each tensor that was split, in the order they were given.

index property

index: tuple[slice, ...]

The index of the chunk's real positions in a tensor of the whole sequence, e.g. output[chunk.index] = ....

positions instance-attribute

positions: slice

The slice of the sequence's real positions that the chunk covers. It is shorter than the chunk size only for a padded last chunk.

tensors instance-attribute

tensors: tuple[Tensor, ...]

The padded chunk of each tensor that was split, in the order they were given.

unpad

unpad(tensor: Tensor) -> torch.Tensor

Drop the padded positions from a tensor computed for this chunk.

Parameters:

Name	Type	Description	Default
tensor	Tensor	A tensor of shape (batch, chunk_size, ...) computed for this chunk.	required

Returns:

Type	Description
torch.Tensor	A view of tensor with only the chunk's real positions, of shape (batch, num_real_positions, ...).

derive_chunk_seed

derive_chunk_seed(seed: int, chunk_index: int) -> int

Derive the seed of the generator that samples the noise of one chunk of a chunked forward pass.

The derived seed depends on chunk_index as well as seed. Seeding every chunk with seed alone would draw the same values for every chunk, making the noise periodic with period chunk_size: anyone who learned the noise of one chunk would know it for all of them.

The derived seed is a hash of seed and chunk_index, so each chunk of each seed gets an unrelated generator. Combining them arithmetically instead, e.g. seed + chunk_index, would give different seeds overlapping chunks, such as chunk 1 of seed 0 and chunk 0 of seed 1, so consecutive per-request seeds would share noise.

Parameters:

Name	Type	Description	Default
seed	int	The noise layer's seed. Values outside the range of a torch.Generator seed are reduced modulo 2**64, which is how a generator stores them.	required
chunk_index	int	The index of the chunk along the sequence dimension, in [0, 2**64).	required

Returns:

Type	Description
int	A seed in [0, 2**64) for the chunk's generator.

Raises:

Type	Description
ValueError	If chunk_index is outside [0, 2**64).

Added in version v4.1.0.

padded_chunks

padded_chunks(
    chunk_size: int, *tensors: Tensor
) -> Iterator[PaddedChunk]

Split tensors along their sequence dimension into consecutive chunks, each padded at the end to chunk_size positions.

Every chunk has the same shape, so a computation run chunk by chunk sees the same shapes whatever the sequence length. Only the last chunk can be padded, with zeros, or False for boolean tensors.

Parameters:

Name	Type	Description	Default
chunk_size	int	The number of positions per chunk.	required
*tensors	Tensor	The tensors to split, each of shape (batch, sequence, ...) with the same sequence length.	required

Yields:

Type	Description
Iterator[PaddedChunk]	Each chunk, with the slice of the sequence's real positions it covers and the padded chunk of each tensor, in order.

Raises:

Type	Description
ValueError	If chunk_size is not positive.
ValueError	If the tensors do not all have the same sequence length.

Added in version v4.3.0.

sample_chunked_noise

sample_chunked_noise(
    shape: Sequence[int],
    *,
    seed: int,
    chunk_size: int,
    device: device,
    dtype: dtype
) -> torch.Tensor

Sample standard normal noise whose values at each position depend only on seed and the position's chunk.

The sequence dimension is split into chunks of chunk_size positions. Each chunk is drawn in full, with shape (1, chunk_size, ...), from a generator freshly seeded by derive_chunk_seed(seed, chunk_index), broadcast to every batch row, and the last chunk is trimmed to the sequence length. Every draw therefore has the same shape regardless of the sequence length and the batch size, so the noise of the first n positions is identical whether n or more positions are sampled, and whichever row a sequence is in. Drawing a chunk for the whole batch at once would instead give each row a different slice of the draw, so a sequence's noise would change with its row.

Neither the global random number generator nor any persisted generator is advanced, so the noise is also identical across calls.

The noise is drawn per chunk rather than per position because only fixed-shape draws are reproducible, and each draw has a fixed cost. A single draw over the whole sequence is not prefix-stable: which random values land at a position depends on the total shape. A draw per position would be, but each draw costs a generator re-seed and a kernel launch, which made one draw per position about 60 times slower than a single draw for an 8192-token sequence.

Note

With a fixed seed, the noise does not depend on the input: every input sampled with the same seed, chunk_size, feature shape, device type, and dtype receives the same noise at each position, across the whole sequence. That includes every row of a batch, so sequences batched together share their noise, exactly as if each were sampled alone with that seed. The generators are per device type, so the same seed yields different noise on different device types.

Parameters:

Name	Type	Description	Default
shape	Sequence[int]	The shape of the noise, (batch, sequence, ...).	required
seed	int	The seed to derive the per-chunk generator seeds from.	required
chunk_size	int	The number of positions per chunk.	required
device	device	The device to sample the noise on.	required
dtype	dtype	The dtype of the noise.	required

Returns:

Type	Description
torch.Tensor	Standard normal noise of the given shape, device, and dtype.

Raises:

Type	Description
ValueError	If chunk_size is not positive.

Added in version v4.1.0.