Add NonErgodicGenerativeProcess and InflatedVocabularyProcess

[kylejray]

Summary

• Add NonErgodicGenerativeProcess — a block-diagonal mixture model composing multiple GenerativeProcess components with weighted probabilities, supporting HMM, GHMM, and FactoredGenerativeProcess components
• Add InflatedVocabularyProcess — a wrapper that multiplies vocab size by factor K with uniform noise, increasing optimal loss by exactly log(K) nats
• Add builder functions for disjoint and partial-overlap vocabulary configurations with prefix, sliding, and random vocab map modes
• Use IndependentFactoredGenerativeProcess for independent structures (O(sum V_i) vs O(prod V_i) sampling)
• Fix off-by-one in NonErgodic generate inference pass (prior vs posterior states)
• Fix random vocab mode to independently sample tokens per component

Test plan

• 21 unit tests for NonErgodicGenerativeProcess (block-diagonal structure, vocab mapping, generate/loss, multiple component types)
• Tests for InflatedVocabularyProcess (vocab inflation, loss offset verification, state dynamics)
• Tests for builder functions (disjoint/partial-overlap vocab modes, config-driven instantiation)

🤖 Generated with Claude Code

[kylejray]

@claude please review this PR!

[claude]

Claude finished @kylejray's task in 3m 6s —— View job

---

PR Review: Add NonErgodicGenerativeProcess and InflatedVocabularyProcess

Todo List:

• Gather context and understand the PR
• Read changed files
• Analyze code quality, correctness, and test coverage
• Post review feedback

---

Overall this is a well-structured PR with clear semantics and solid test coverage. The core math is correct. I found one meaningful correctness bug, a few design issues, and some style violations to flag.

---

Critical Bug: _slice_belief_states Temporal Misalignment

File: simplexity/generative_processes/generator.py:107-113

When generate(return_all_states=True) is called on a NonErgodicGenerativeProcess, JAX's scan stacks the inference trajectory into a NonErgodicState where both component_beliefs and component_states[i] have the time dimension. After vmap, shapes are:

• component_beliefs: [B, T, num_components]
• component_states[i]: [B, T, state_dim_i]

But _slice_belief_states only slices component_beliefs and leaves component_states unmodified:

if isinstance(belief_states, NonErgodicState):
    return NonErgodicState(
        component_beliefs=belief_states.component_beliefs[:, seq_slice, ...],  # sliced
        component_states=belief_states.component_states,  # ← NOT sliced, still [B, T, ...]
    )

This produces a NonErgodicState with mismatched time dimensions between component_beliefs and component_states whenever generate_data_batch_with_full_history is called. Fix this →

Related: there is no test for NonErgodicGenerativeProcess.generate(return_all_states=True), which is especially important given the off-by-one fix mentioned in the PR description.

---

Missing Validation: Assumed Equal Vocab Sizes in build_nonergodic_partial_overlap

File: simplexity/generative_processes/builder.py:846

v = temp.components[0].vocab_size  # silently ignores other components

If components have different vocab sizes (e.g. a coin and a mess3), n_shared and n_unique are computed only from component 0. For "sliding" and "random" modes this silently produces incorrect vocab maps. Should validate that all components share the same vocab_size before proceeding, or at minimum document the constraint. Fix this →

---

Inline Comments Violate CLAUDE.md Style Guidelines

File: simplexity/generative_processes/nonergodic_generative_process.py:17-37 and throughout generate()

CLAUDE.md states: "No Comments: Avoid inline comments; code should be self-documenting through clear naming and structure."

The file contains a large module-level comment block (lines 17–37) explaining the flatten/pad/unflatten strategy, plus extensive inline comments inside generate(). The explanatory content is valuable but should live in the class/method docstrings rather than as free-floating comment blocks.

---

Double Build in build_nonergodic_disjoint_vocab

File: simplexity/generative_processes/builder.py:765-774

The function builds the entire process once to query vocab sizes, then rebuilds it with the disjoint maps — creating all JAX arrays (transition matrices, inverse maps, etc.) twice. A cleaner pattern would be to extract a helper that builds only the component processes, queries their vocab_size, computes the maps, and then wraps once. Since build_nonergodic_process_from_spec is called twice this is a concrete inefficiency.

---

Missing Tests

1. return_all_states=True for NonErgodicGenerativeProcess — No test exercises the inference trajectory path, which is precisely where the off-by-one fix lives. This is the most important missing test.

2. FactoredGenerativeProcess as a component — The PR explicitly supports factored component types and includes the path in build_nonergodic_process_from_spec, but there is no test for this combination (which uses the flatten/pad/unflatten logic for heterogeneous state types).

3. generate_data_batch_with_full_history with NonErgodicGenerativeProcess — Would catch the _slice_belief_states bug described above.

---

Minor Notes

• build_inflated_process (builder.py:864-877) is a one-line passthrough to InflatedVocabularyProcess(base_process, inflation_factor) with no additional logic. Per CLAUDE.md ("Don't create helpers, utilities, or abstractions for one-time operations"), this wrapper likely doesn't justify its existence.

• YAML config fields (nonergodic_example.yaml:7-8): base_vocab_size: ??? and vocab_size: ??? are Hydra "required" markers. Since these fields aren't consumed by the builder function, any Hydra instantiation attempt will error unless these are explicitly provided or removed.

• NonErgodicState.component_states type is tuple[Any, ...]. A union type tuple[jax.Array | tuple[jax.Array, ...], ...] would be more precise and improve IDE/pyright support.

---

What Looks Good

• The Bayesian filtering in transition_states is mathematically correct, including the safe-indexing with jnp.clip and the 0-likelihood guard for Bayes update.
• The lax.switch + flatten/pad/unflatten approach for heterogeneous state types in generate is well-reasoned and the comments in the generate docstring explain it clearly.
• InflatedVocabularyProcess is clean and the log(K) loss increase property is correctly implemented and tested.
• The vocab map builder tests are thorough, covering disjoint, prefix, sliding, and random modes with seed reproducibility checks.
• IndependentFactoredGenerativeProcess is correctly used for independent structures (O(sum V_i) vs O(prod V_i)).

[kylejray]

@claude can you take another look?

[claude]

Claude finished @kylejray's task in 2m 15s —— View job

---

Re-reviewed PR #172 — All critical and major issues from the previous review have been addressed. Two minor items remain.

Todo List

• Gather context
• Read changed files to assess what's been addressed
• Compare against previous review findings
• Post updated review

---

What Was Fixed

Critical Bug ✅ — _slice_belief_states now correctly slices component_states along axis 1 alongside component_beliefs, and handles nested tuples (FactoredState) via the inner isinstance(cs, tuple) branch in generator.py:107-114.

Missing Validation ✅ — build_nonergodic_partial_overlap now validates equal vocab sizes at builder.py:873-877 before computing n_shared/n_unique.

Inline Comments ✅ — The module-level comment block is gone; the implementation rationale lives in docstrings on _unflatten_state and generate().

Double Build ✅ — build_nonergodic_disjoint_vocab now uses _build_components_from_spec (a new shared helper) so components are built once at builder.py:789.

Missing Tests ✅ — Three new test classes address all gaps:

• TestGenerateReturnAllStates — verifies shapes and belief validity for return_all_states=True
• TestFactoredComponent — exercises the flatten/pad/unflatten path with heterogeneous state types
• TestGenerateDataBatchWithFullHistory — catches _slice_belief_states misalignment, including with BOS token

Type annotation ✅ — component_states is now tuple[ComponentState, ...] where ComponentState = jax.Array | tuple[jax.Array, ...], which is much more precise than tuple[Any, ...].

---

Remaining Items

1. YAML config: spurious required fields

tests/end_to_end/configs/generative_process/nonergodic_example.yaml:7-8

base_vocab_size: ???  # Will auto-infer as 3
vocab_size: ???

These top-level ??? fields exist outside the instance: block, so they are never passed to build_nonergodic_process_from_spec. However, Hydra treats ??? as "mandatory" — any attempt to load this config without explicitly overriding these keys will throw MissingMandatoryValue. The comment # Will auto-infer as 3 is also incorrect (Hydra cannot infer them). These should either be given concrete values, set to null, or removed entirely. Fix this →

2. build_inflated_process passthrough (minor, judgment call)

builder.py:900-913 is still a one-liner wrapping InflatedVocabularyProcess(base_process, inflation_factor) with no additional logic. Per CLAUDE.md: "Don't create helpers, utilities, or abstractions for one-time operations." build_inflated_process_from_spec is the useful function here (it parses a spec dict); build_inflated_process adds nothing over calling the constructor directly. Removing it keeps the API surface minimal. Either way this is minor and doesn't block the PR.

---

Overall

The PR is in good shape. The critical correctness fix and all the missing tests are now present. The YAML issue is the only thing I'd call out before merging.
| Branch: kyle/nonergodic

[kylejray]

@casperlchristensen

I think we are in good shape now, as far as the tests are concerned.

Re: claude's comments:

• I think the ??? in the yaml is fine and consistent with what we do.
• I think the inflated_process "passthrough" also is there to preserve the interface pattern of other processes builders. its a trivial helper, but maybe makes sense here. Not sure on that one.

[casperlchristensen]

@claude Please re-check post fixing merge conflicts!

[claude]

Claude finished @casperlchristensen's task in 4m 16s —— View job

---

Post-Merge Re-check

Todo List:

• Gather context from previous reviews
• Inspect current diff vs main
• Read key changed files
• Check merge conflict resolution for correctness
• Post updated review

---

CI status: ✅ Both workflows passing (End-to-End Tests + main workflow) on the post-merge commits.

---

Merge Conflict Resolution

Three files had merge conflicts. All are resolved correctly:

simplexity/generative_processes/structures/fully_conditional.py
Main merged in a new indexing.py module with build_other_factor_multipliers and flatten_index. The conflict resolution correctly replaced the old jnp.sum(tokens * mult) call (with an inline comment) with flatten_index(tokens, mult) from the new shared module — equivalent logic, better factored. ✅

simplexity/utils/factoring_utils.py
Main added _radix_multipliers, compute_other_multipliers, and compute_prefix_multipliers to this file. The PR added docstrings to compute_other_multipliers and compute_prefix_multipliers. Conflict resolved by correctly merging both sets of changes. ✅

simplexity/generative_processes/structures/conditional_transitions.py
PR's docstring additions cleanly merged with main's changes. ✅

---

Previously Flagged Items — Status

All critical and major items from the two prior reviews remain resolved. The two minor open items from the second review:

• YAML ??? fields — The misleading # Will auto-infer as 3 comment was removed in commit 2520382. The ??? markers remain (PR author's call to keep them consistent with existing configs). ✅ addressed.
• build_inflated_process passthrough — Still present; PR author keeps it for API consistency. Minor, judgment call. No regression.

---

New Observation: assert seed is not None (minor, correct)

builder.py:887 has assert seed is not None inside the elif mode == "random": branch, after the earlier guard:

if mode == "random" and seed is None:
    raise ValueError("seed is required when mode='random'")

The assert is retained to narrow seed: int | None for pyright, since the combined condition in the earlier guard doesn't allow pyright to infer seed is not None at the elif branch. This is idiomatic and correct — no issue.

---

Scope Note

The PR title says "Add NonErgodicGenerativeProcess and InflatedVocabularyProcess" but the diff includes a large new visualization system that wasn't mentioned in the title or description:

• simplexity/visualization/ — 6 new files (~2,500 lines including a 1,338-line plotly_renderer.py)
• simplexity/activations/visualization/ — 8 new files (~2,100 lines)

The previous reviews did not cover this code. The associated tests pass in CI. If the visualization system is intended to be part of this PR, it may warrant a dedicated description section or a separate PR for easier review.

---

Overall

The merge conflict resolution is clean — no correctness regressions introduced. All previously flagged critical issues remain fixed. The PR is in the same good state as after the second review. The only new consideration is whether the visualization code should be reviewed/scoped separately, which is a process question rather than a correctness concern.
| Branch: kyle/nonergodic

[casperlchristensen]

@claude The visualization was not meant to be added -- it has been re-deleted!

[claude]

Claude finished @casperlchristensen's task in 33s —— View job

---

Post-Merge Re-check (visualization removed)

Todo List:

• Gather context from previous reviews
• Check current diff vs main to confirm visualization is removed
• Review remaining changed files for any new issues
• Post updated review

---

Visualization removed ✅ — The re-delete visualization commit (e67ab10) correctly removes all visualization files. The current diff vs origin/main contains exactly the 16 files listed in the PR description — no more, no less.

CI status: Both workflows are currently in progress (22648642063, 22648642054). Will need to pass before merging.

---

Current State

All issues from the previous three reviews remain resolved. There are no new findings introduced by the visualization deletion. The PR is scoped to exactly what the description says:

• NonErgodicGenerativeProcess + InflatedVocabularyProcess + builder functions
• Associated tests (21 unit tests for NonErgodic, tests for InflatedVocabulary, builder tests)
• Merge-conflict resolution in fully_conditional.py, factoring_utils.py, conditional_transitions.py (all clean)
• run_parallel.py and pyproject.toml changes

The two minor items from the second review (??? YAML fields, build_inflated_process passthrough) were triaged as intentional by the PR author — no action needed.

The PR looks good to merge once CI passes.

---

| Branch: kyle/nonergodic

[casperlchristensen]

@ealt I'm happy with this now!