forking
forking
¶
DiversityScratchbubbles
¶
Bases: Scratchbubbles
Scratchbubbles training model with fork-diversity instrumentation.
measure_friend_diversity(x: jax.Array, token_index: jax.Array) -> jax.Array
staticmethod
¶
Measure mean cosine distance between non-rightmost siblings.
decode(x: jax.Array, padding_mask: Optional[jax.Array] = None, deterministic: bool = False, **kwargs: Any) -> Any
¶
Decode while recording states immediately after each fork.
ForkAnalysis(spec: ExecutionSpec)
¶
Bases: InferenceJob[ForkAnalysisConfig, Thoughtbubbles]
Export tokenwise fork diversity and a per-fork logit lens.
capture_fork_blocks(module: Any, method: str) -> bool
staticmethod
¶
Capture the state tuple returned by each forking block.
measure_tokenwise_diversity(states: np.ndarray, token_index: np.ndarray, sequence_length: int) -> list[dict[str, int | float | None]]
staticmethod
¶
Measure non-rightmost sibling cosine distance for each input token.