looped

looped

LoopedGPT

Bases: GPT

GPT whose ordinary block stack is applied repeatedly.

n_repeats: int property

Number of applications implied by effective and parameter depth.

new_block() -> Any

Construct one uniquely parameterized block.

apply_block(block: Any, x: jax.Array, repeat: int, padding_mask: Optional[jax.Array], deterministic: bool) -> jax.Array

Apply a block without conditioning it on the repeat index.

HintedLoopedGPT

Bases: LoopedGPT

Looped GPT whose blocks receive the current repeat as a learned hint.

new_block() -> Any

Construct one uniquely parameterized hinted block.

apply_block(block: Any, x: jax.Array, repeat: int, padding_mask: Optional[jax.Array], deterministic: bool) -> jax.Array

Apply a block with the current repeat as its hint index.