Index
block
¶
Block
¶
ThoughtBlock
¶
Bases: Block
Transformer block for Thoughtbubbles with cumulative score support. Extends Block by weighting attention and MLP outputs by cumulative scores.
__call__(x: jax.Array, **kwargs: Any) -> Union[jax.Array, ForkingOutput]
¶
Forward pass with cumulative score weighting.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Array
|
Input tensor of shape (B, T, C). |
required |
**kwargs
|
Any
|
Must include cumulative_scores, token_index for forking mode. |
{}
|
Returns:
| Type | Description |
|---|---|
Union[Array, ForkingOutput]
|
If forking: Tuple of (output, cumulative_scores, token_index). |
Union[Array, ForkingOutput]
|
If not forking: Just the output tensor. |
ForkingBlock
¶
Bases: ThoughtBlock
Transformer block with forking capability. Extends ThoughtBlock by adding token forking based on learned scores.
clipped_logsigmoid(x: jax.Array, min_val: float = -20.0) -> jax.Array
staticmethod
¶
Compute log-sigmoid with clipping for numerical stability.
fork(x: jax.Array, cumulative_scores: jax.Array, token_index: jax.Array, padding_mask: Optional[jax.Array], input_seq_len: Optional[int] = None) -> Tuple[jax.Array, jax.Array, jax.Array]
¶
Top-k forking: doubles tokens then selects top-k to keep.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Array
|
Input tensor of shape (B, T, C). |
required |
cumulative_scores
|
Array
|
Log-probability scores of shape (B, T). |
required |
token_index
|
Array
|
Original token indices of shape (B, T). |
required |
padding_mask
|
Optional[Array]
|
Boolean tensor of shape (B, T). True for valid. |
required |
input_seq_len
|
Optional[int]
|
Original input sequence length for ratio calc. |
None
|
Returns:
| Type | Description |
|---|---|
Tuple[Array, Array, Array]
|
Tuple of (forked_x, new_scores, new_token_indices). |
fork_token(x: jax.Array, parent_indices: jax.Array, *, cumulative_scores: jax.Array, token_index: jax.Array, padding_mask: Optional[jax.Array]) -> jax.Array
¶
Return candidate replacements at parent_indices from the full source x.
Source scores, original positions, and padding support contextual forks. The caller applies replacements only to new children, keeping selection fixed-size for JAX and retaining keep children unchanged.
__call__(x: jax.Array, **kwargs: Any) -> Union[jax.Array, ForkingOutput]
¶
Forward pass: fork first, then apply standard block operations.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Array
|
Input tensor of shape (B, T, C). |
required |
**kwargs
|
Any
|
cumulative_scores, token_index, padding_mask, deterministic, input_seq_len. |
{}
|
Returns:
| Type | Description |
|---|---|
Union[Array, ForkingOutput]
|
Tuple of (output, cumulative_scores, token_index). |