diff --git a/docs/design/generators_instantiation.md b/docs/design/generators_instantiation.md new file mode 100644 index 00000000..efcfe197 --- /dev/null +++ b/docs/design/generators_instantiation.md @@ -0,0 +1,616 @@ +# Instantiating vendored generative processes under `managed_run` + +**Status:** accepted and implemented; see §8 for the decisions taken and §9 for what remains open +**Scope:** how `simplexity`'s run management instantiates generative processes that live in the +*consumer's* repository, as a precondition for deprecating `simplexity/generative_processes` +in favour of [`generators`](https://github.com/ealt/generators) +**Branch:** `feat/generators-instantiation` + +--- + +## 1 Summary + +`generators` distributes generative-process code by **copying**: consumers vendor the modules into +their own project and modify them freely. `simplexity`'s runner instantiates components from Hydra +config. Those two facts are compatible. What is *not* compatible is one specific mechanism: +`simplexity` decides whether a config section is a generative process by testing whether its +`_target_` string starts with `"simplexity.generative_processes."`. Vendored code can never satisfy +a namespace test, because refusing to be a shared namespace is the entire point of the vendoring +model. + +So the tension is narrower than "config-driven instantiation vs. a library that refuses to be a +dependency." Hydra instantiating consumer-local code is not a problem — it already happens for +`torch.optim.Adam` and `HookedTransformer` on every run. **The problem is that +`generative_process` is the one component whose discovery predicate is a namespace prefix and whose +instantiation asserts a `simplexity` base class.** Fixing it means replacing a *nominal* test +(where does this code live?) with a *structural* one (what can this object do?) — which is what +`generators`' own SPEC licenses: + +> §2 Operations — "The spec defines what results are required, not how they are organized or +> computed." + +**Recommendation, in one line:** make the process contract a `runtime_checkable` Protocol keyed on +the SPEC's operations, let the config section *declare* itself as a generative process instead of +being sniffed by namespace, and make a declared-but-nonconforming section a hard error instead of a +silent skip. Deprecate `simplexity/generative_processes` by warning and documentation, not +deletion, and only after the four capability gaps in §4.3 have homes. + +This was accepted and is implemented on this branch — §8 records the decisions, §9 what is still +open, §10 what changed. §§2-7 are the analysis as written before the decision, kept as the argument +for it. `docs/generators_migration.md` is the consumer-facing guide. + +Measured up front, so the rest of this document argues from evidence rather than expectation +(§3 has the reproduction): + +| Question | Measured answer | +| --- | --- | +| Does a vendored `generators` process work with `simplexity`'s generation path today? | **Yes** — `generate_data_batch` runs on it unmodified, correct shapes, normalized `obs_dist`, correct `seq_prob`. | +| How much of the `GenerativeProcess` interface does it miss? | **2 of 8 members**: `log_observation_probability_distribution`, `log_probability`. Nothing on the training path. | +| What actually blocks it? | `is_generative_process_target` → `False`, so the config key is filtered out before instantiation. | +| What does the user see when that happens? | One `INFO` line, `[generative process] no generative process configs found`. **No warning, no error.** | +| Does a structural Protocol also accept the existing in-repo processes? | **Yes** — same predicate accepts `build_hidden_markov_model`'s output, so this is not a breaking change. | + +--- + +## 2 The tension, precisely + +### 2.1 What the runner does today + +Four places in `simplexity/run_management/run_management.py` participate: + +| Step | Code | Mechanism | +| --- | --- | --- | +| **Discover** which config sections are processes | `_setup_generative_processes` → `filter_instance_keys(..., is_generative_process_target, ...)` (`run_management.py:348`) | `target.startswith("simplexity.generative_processes.")` (`structured_configs/generative_process.py:388`) | +| **Validate** the config | `validate_generative_process_config` (`structured_configs/generative_process.py:401`) | dispatches on the exact `_target_` of five known builders; the `else` branch re-requires the namespace prefix (line 426) | +| **Instantiate** | `_instantiate_generative_process` (`run_management.py:330`) | `typed_instantiate(instance_config, GenerativeProcess)` → `assert isinstance(obj, GenerativeProcess)` (`utils/config_utils.py:145`) | +| **Resolve vocab metadata** | `_get_attribute_value` (`run_management.py:409`) | re-runs the *same* namespace filter to find `vocab_size` / `bos_token` for the predictive model | + +Note the fourth row: the namespace filter is not only how the process is built, it is also how +`vocab_size` reaches the predictive-model config (`_setup_predictive_models`, `run_management.py:470`). +A process the filter rejects therefore fails **twice** — no process, and no vocab-size resolution. + +### 2.2 Why `generators` cannot satisfy it + +`generators/USAGE.md` opens: + +> This repo is not a library. Do not add it to your dependencies. Copy the modules you need into +> your codebase. + +and states the mechanism explicitly: "The unit of reuse is the module, and the mechanism is +copying… The invariant that survives your modifications is not the code — it is the spec." The +repo has empty `__init__.py` files and no registry "coupling modules together" on purpose. + +A vendored process therefore has: + +- a `_target_` under the *consumer's* package name, never `simplexity.*`; +- a shape the consumer has deliberately changed (deletion is the stated specialization mechanism); +- no class at all — `generators` is **module-level functions over a `NamedTuple`**, e.g. + `generators/ghmm/process.py` exposes `init(Ts) -> Data`, `obs_dist(data, eta)`, + `sample(data, eta, key)`, `update(data, eta, x)`, `generate(data, eta, keys)`, + `seq_prob(data, xs)`. There is nothing for `isinstance(obj, GenerativeProcess)` to be true of. + +### 2.3 The reframe + +`simplexity` uses a namespace prefix as a proxy for a capability. That proxy is *sound* while all +processes live in `simplexity`, and *unfixable* once they deliberately do not. Every option below +is a different answer to "what replaces the proxy?" + +Worth noting: **`simplexity` has already made this move for two other components.** + +- `is_predictive_model_target` (`structured_configs/predictive_model.py:192`) accepts any + `*.nn.*` or `*.models*` target — `torch.nn`, `equinox.nn`, `penzai.models`. +- `is_optimizer_target` (`structured_configs/optimizer.py:77`) accepts `torch.optim.*` and `optax.*`. +- `_instantiate_predictive_model` (`run_management.py:437`) calls bare `hydra.utils.instantiate` + with **no** type assertion at all — the predictive-model contract is already fully structural. + +So relaxing the generative-process contract is not a new pattern in this repo. It is bringing the +most rigid component filter into line with how the repo already treats the two components whose +implementations live upstream. (Both existing relaxations are still *namespace sniffing*, just with +a wider net — `parts[1] == "nn"` would not accept a consumer-local module either. §5.2 argues that +sniffing is the wrong axis regardless.) + +--- + +## 3 The measured interface delta + +The brief asked for the actual delta rather than an assumed one. Reproduction: vendor +`generators/ghmm/process.py` + `generators/utils.py` + `transition_matrices/classical.py` into a +throwaway package, write a ~50-line adapter binding the module functions to method names, and run +it through `simplexity`'s real code paths. + +### 3.1 Operation-by-operation mapping + +| `GenerativeProcess` member | `generators` (ghmm) | Delta | +| --- | --- | --- | +| `vocab_size` | `data.Ts.shape[0]` | derived, trivial | +| `initial_state` | `data.eta_0` | field on `Data` | +| `emit_observation(state, key)` | `sample(data, eta, key)` | same math, `data` threaded as an argument | +| `transition_states(state, obs)` | `update(data, eta, x)` | same | +| `observation_probability_distribution(state)` | `obs_dist(data, eta)` | same | +| `probability(observations)` | `seq_prob(data, xs)` | same for ghmm; **signature diverges** for factored (`seq_prob(data, eta, xs, *, decode)`) | +| `generate(state, key, len, return_all)` | `generate(data, eta, keys)` | `generators` returns final state only; `return_all_states=True` needs a 3-line scan change. Batching is the consumer's (`eqx.filter_vmap` vs. explicit `vmap`) | +| `log_observation_probability_distribution` | — | **absent** | +| `log_probability` | — | **absent** | + +Two genuine gaps out of eight members, and both are log-space. Their call sites are worth +knowing before deciding whether the contract must require them: + +- `log_observation_probability_distribution` is consumed by `mixed_state_presentation.py:512` + (the mixed-state-tree / belief-enumeration machinery) and by the composite processes that + forward to sub-processes (`nonergodic_generative_process.py:254`, + `inflated_vocabulary_process.py:88`, `factored_generative_process.py:148`). +- Neither log-space method is called anywhere on the **training or generation path** — + not by `generator.py`, `torch_generator.py`, `run_management.py`, or `tests/end_to_end/training.py`. + +That asymmetry is the evidence for a **tiered contract** (§5.3): the operations training needs and +the operations belief-analysis needs are different sets, and today's ABC forces every process to +implement both. + +### 3.2 The blocking mechanism, measured + +Running the vendored adapter through the real predicates: + +``` +1. isinstance(vendored, GenerativeProcess) = False +2. is_generative_process_target('vendored.adapter.build_vendored_mess3') = False +3. instance keys discovered = ['generative_process.instance'] + keys surviving the generative-process filter = [] + -> components.generative_processes would be None (SILENT skip) + validate_generative_process_config raised ConfigValidationError: + GenerativeProcessConfig.instance must be a generative process target +5. GenerativeProcess abstract members = ['emit_observation', 'initial_state', + 'log_observation_probability_distribution', 'log_probability', + 'observation_probability_distribution', 'probability', 'transition_states', 'vocab_size'] + missing on the vendored adapter = ['log_observation_probability_distribution', 'log_probability'] +4. generate_data_batch on the vendored process: inputs (4, 16), labels (4, 16) + vocab_size=3 token range=[0,3] + obs_dist(initial) = [0.33333325 0.3333333 0.3333333 ] sums to 1: True + probability(xs) = 0.0028251 +``` + +Read line 4 first: **the vendored process already works.** `generate_data_batch` — the function +`tests/end_to_end/training.py` and three of the four external consumers call — runs on it +unmodified and produces correct results. No change to `simplexity`'s *runtime* code is required +for a vendored process to train. The work is entirely in discovery and validation. + +### 3.3 The failure mode is silent + +`filter_instance_keys` (`utils/config_utils.py:67`) short-circuits: + +```python +if isinstance(target, str) and filter_fn(target) and _validate(cfg, instance_key, validate_fn, ...): +``` + +Because `filter_fn` returns `False` first, `_validate` is never called, so the +`ConfigValidationError` above is **never raised in the real path** and its warning is never logged. +Confirmed by driving `_setup_generative_processes` directly at `DEBUG`; the complete output is: + +``` +INFO simplexity: [generative process] no generative process configs found +RESULT: None +``` + +A config that names a vendored process is indistinguishable from a config that has no process at +all. `managed_run` proceeds, `components.get_generative_process()` returns `None`, and the run +fails somewhere downstream — or, if the entrypoint is tolerant, does something wrong quietly. +**This is a defect independent of which option below is chosen, and worth fixing either way:** a +config section typed as `GenerativeProcessConfig` that yields no process should be loud. + +### 3.4 Assembly is not expressible in YAML + +`generators`' composite modules take **keyword-only callables** the consumer must supply: + +```python +def obs_dist(data, eta, *, decode: Callable[[jax.Array], jax.Array]) -> jax.Array +def generate(data, eta, keys, *, encode: Callable[[jax.Array], jax.Array]) -> tuple[...] +``` + +(`generators/factored/chain.py:106,164`; same in `complete.py`, `independent.py`, +`nonergodic/factored.py`.) The composite-token encoding is deliberately not owned by the module — +SPEC §4.3.2–4.3.3 specify radix encoding but the module takes it as an argument. This is what +`ZEN.md`'s "Some assembly required" means concretely. + +**Consequence for config design:** a vendored process cannot generally be described by a `_target_` +plus scalar kwargs, because part of its definition is *code the consumer writes*. Whatever contract +we choose, the config for a vendored process should point at a **project-local factory** that does +the assembly and returns the finished object. That is not a new mechanism — `simplexity`'s own +`_target_: simplexity.generative_processes.builder.build_hidden_markov_model` is exactly this +pattern, and it is what the spike used. + +--- + +## 4 The strategic fork (please resolve first) + +"Deprecate `generative_processes` in favour of `generators`" has two readings, and they imply +different work. Everything downstream depends on which one is intended. + +### 4.1 Reading A — `simplexity` becomes a runner, not a process home + +`simplexity` stops being where processes live. It keeps the runner, the training/eval/analysis +machinery, and a *contract*; processes live in each consumer's repo as vendored `generators` +modules. `simplexity/generative_processes` gains deprecation warnings, stops growing, and +eventually thins to the protocol plus the consumer-side utilities that operate on it +(`generator.py`, `torch_generator.py`, `mixed_state_presentation.py`). + +This is what the brief describes and what `generators`' USAGE.md argues for — the coordination +point disappears; research directions decouple. + +### 4.2 Reading B — `simplexity` vendors `generators` internally + +`simplexity` copies `generators`' modules into `simplexity/generative_processes` and reimplements +its process classes on top of them, so in-repo processes become spec-conformant. + +This is worth naming because it is the tempting middle path, and it **defeats the purpose for +downstream consumers**: `simplexity` would still be the shared library that everyone imports +processes from, still be the merge queue, still be frozen by the reproducibility obligation. It +buys spec conformance for `simplexity`'s own copies and nothing else. It is not obviously wrong — +if the goal is only "adopt `generators`' math," it is cheaper — but it does not deliver the +decoupling. + +**Recommendation: Reading A.** The rest of this document assumes it. + +### 4.3 What Reading A owes: `generators` is not a superset + +Deprecation implies a replacement exists. Measured against `simplexity/generative_processes` +(3,754 lines across 21 modules), `generators` does not yet cover four things. These are not +objections — they are work that needs an owner before any removal: + +1. **Mixed-state presentation / belief enumeration.** `mixed_state_presentation.py` (576 lines) — + `MixedStateTree`, `MixedStateTreeGenerator`, myopic-entropy computation. `generators` has no + counterpart (SPEC covers belief *updates*, not tree enumeration). This is load-bearing for + the interpretability work: it is what produces the ground-truth belief geometry that + `activation_tracker` compares activations against, and `eden-experiments` imports it directly + (`fraxl-run/salvage-candidate/scratch/sweep.py:17`, `nonergodic-demo/nonergodic_entropy.py:106`). +2. **Log-space operations.** The two missing members of §3.1. Required by the MSP path for + numerical stability over long sequences. +3. **Framework bridges and augmentation.** `torch_generator.py` (torch tensors / device + placement), `data_prefetcher.py`, and BOS/EOS handling in `generator.py`. SPEC §6.2 specifies + augmentation behaviorally and Appendix A marks device/DLPack interop explicitly *optional* and + out of conformance scope, so this is consumer-side by design — but it currently lives in the + package being deprecated. +4. **Process zoo breadth.** `transition_matrices.py` has 16 process families (`rrxor`, `fanizza`, + `tom_quantum`, `post_quantum`, `days_of_week`, `even_ones`, `matching_parens`, + `no_consecutive_ones`, `mr_name`, `sns`, `coin`, `leaky_rrxor`, …). `generators`' + `transition_matrices/` has `zero_one_random`, `cycle`, `mess`, `bloch_walk`, `moon`, plus + `expand_vocab`/`compress_vocab`. Note these are *matrices* — data, not behavior — and + `generators`' USAGE.md says "Process definitions are your data, not our code… the repo does not + own a process zoo," yet it ships `transition_matrices/` anyway. The status of this directory is + genuinely ambiguous and it is the most reusable, least controversial part of what would be + deprecated (see open question Q5). + +There is also a **missing safety net**: SPEC §7 defines the JSON encoding for conformance test +vectors, and USAGE.md's consumption recipe step 2 is "copy the relevant conformance vectors and +wire them into your test suite." **No vector files exist in the `generators` repo yet** — no JSON +fixtures, no test referencing them. The invariant that is supposed to survive a consumer's +modifications is currently unenforceable. That matters directly here: a structural contract is +only as safe as the consumer's ability to prove their fork still implements the same math. + +--- + +## 5 Options + +Three independent axes. The brief's three options map onto them; presenting them as axes rather +than as three rival packages avoids a false choice, because the brief's option 3 +(factory/callable config) is not an alternative to options 1 and 2 — it is the config *shape* +either of them needs (§3.4). + +### 5.1 Axis A — what is the contract? + +| | **A1: keep the nominal ABC** | **A2: `runtime_checkable` Protocol** | **A3: no check (predictive-model style)** | +| --- | --- | --- | --- | +| Mechanism | consumer subclasses `simplexity.generative_processes.generative_process.GenerativeProcess` | consumer's object structurally satisfies a Protocol; verified at instantiation | bare `hydra.utils.instantiate`, duck-typed at use site | +| Hydra ergonomics | unchanged | unchanged | unchanged | +| Type checking | strongest (nominal) | good — pyright checks conformance statically *if* the consumer annotates | none | +| Preserves `generators`' point? | **No** — requires importing the deprecated package and inheriting from it; recreates the shared coordination surface | **Yes** — a Protocol is a type-only dependency; nothing behavioral is shared, consumer may still delete/rename/restructure freely | Yes, maximally | +| Failure mode | clear (`TypeError` on abstract methods) | clear if we raise (`isinstance` False → explicit error) | **poor** — `AttributeError` deep inside a `jit`-traced scan, far from the config that caused it | +| Reproducibility asserts | unaffected | unaffected | unaffected | +| Verified? | n/a | **yes** — accepts vendored adapter *and* native `HiddenMarkovModel`, rejects a bogus object (§3.2, §5.5) | n/a | + +A1 is disqualified by the requirement, not by taste: inheriting from a class inside the package +being deprecated is a behavioral dependency on that package. A3 is what the repo already does for +predictive models, so it is defensible on consistency grounds, but it trades a config-time error +for a runtime one — and given that the current failure is *already* an invisible `None` (§3.3), +adding more silence is the wrong direction. + +**Recommend A2.** + +A Protocol's honest limits, since they should be stated rather than discovered later: +`isinstance` against a `runtime_checkable` Protocol checks **member presence only** — not +signatures, not shapes, not semantics. It will not catch a `transition_states` with the arguments +reversed, a `vocab_size` that lies, or an `obs_dist` that returns unnormalized values. It is a +smoke test that replaces a *namespace* check, not a correctness proof. What actually protects +correctness is the SPEC's conformance vectors run against the consumer's copy — which is why their +absence (§4.3) is a real risk and Q3 below matters. + +### 5.2 Axis B — how does the runner discover the component? + +The `_target_` string is the wrong discriminator for consumer-local code, and no widening of the +prefix test fixes that: any prefix wide enough to admit arbitrary consumer packages admits +everything, including the logger and the optimizer. + +| | **B1: widen the prefix** | **B2: declared in config** | **B3: prefix fast path + explicit declaration** | +| --- | --- | --- | --- | +| Mechanism | add more accepted prefixes | the config section declares its kind; namespace ignored | in-repo targets keep the prefix path; consumer-local sections opt in with an explicit marker | +| Works for arbitrary consumer packages? | no | yes | yes | +| Existing configs change? | no | **yes** — every process config needs the new field | no | +| Risk | unbounded — a wide prefix mis-claims other components | migration churn across 11 in-repo configs + teammates' branches | small: two paths to maintain | + +B2 is the clean end-state and it is *already half-true*: `_instantiate_generative_process` and +`_get_attribute_value` both derive `config_key = instance_key.rsplit(".", 1)[0]` and assume the +parent section is the process config block, and that section is already typed as +`GenerativeProcessConfig` in every entrypoint's structured config +(`tests/end_to_end/training.py:66`). The target-string filter is largely **redundant with the +structured-config schema** that already declares intent. + +**Recommend B3** — B2's mechanism, added without breaking the 11 existing in-repo configs or +teammates' in-flight branches. Concretely: `is_generative_process_target` keeps returning `True` +for `simplexity.generative_processes.*`, and the section may additionally opt in via an explicit +field on `GenerativeProcessConfig`. A section that opts in and then fails the Protocol check is a +**hard error**, not a filtered key — that is the §3.3 fix. + +### 5.3 Axis C — what surface does the contract require? + +Today's ABC demands all eight members from every process. §3.1 measured that the training path +uses six and the belief-analysis path uses the other two. + +| | **C1: all eight** | **C2: core + optional extensions** | +| --- | --- | --- | +| Required of a vendored process | 8 members incl. log-space | 6 core; log-space only if the run uses MSP/analysis | +| Cost to consumer | must hand-write log-space ops `generators` does not provide | copies what the SPEC covers; adds log-space when analysis needs it | +| Enforcement | one `isinstance` | core checked at instantiation; extension checked where used (MSP generator, composite processes) | +| Honesty | over-claims — asserts capabilities training never exercises | matches measured usage | + +**Recommend C2**, with the core protocol keyed on SPEC §2 + §6.1 (obs dist, sampling, state update, +sequence probability, generation) plus the two metadata members the runner itself needs +(`vocab_size` for vocab resolution, `initial_state` for batch expansion). This also makes the +deprecation *smaller*: log-space is an analysis concern, so it belongs with the analysis code that +consumes it, not in the base contract. + +### 5.4 Recommended shape, concretely + +1. **New protocol module outside the deprecated package** — the contract must not live in + `simplexity/generative_processes/`, or consumers still import from the thing being deprecated. + Suggested home: `simplexity/run_management/` (alongside `components.py`, which is what + type-annotates the object) or a new top-level `simplexity/processes/`. Naming to be settled; + `Protocol` is already an established idiom in this repo + (`generative_processes/structures/protocol.py:39`). +2. **`GenerativeProcess` (the ABC) structurally satisfies the new Protocol** — so every existing + process passes unchanged, and `Components.generative_processes` retypes to the Protocol with no + change to in-repo processes or their configs. Verified: §5.5. +3. **Discovery** via B3: prefix fast path, plus explicit opt-in for consumer-local sections. +4. **Validation** raises on a declared-but-nonconforming section, naming the missing members. + Replaces the silent skip. +5. **Vendored configs point at a project-local factory** (§3.4), documented with a worked example. +6. **Acceptance test:** a vendored `generators` module under `tests/` trains end-to-end under + `@managed_run()` — the brief's real bar. The spike (§3) is that test's skeleton; it needs + promoting from a throwaway to a fixture, with the vendored copy checked in so the test does not + depend on `generators` being present. + +**Incidental cleanups this touches** (flagging rather than silently bundling): + +- `_setup_generative_processes` (`run_management.py:359-365`) calls `_instantiate_generative_process`, + which already calls `resolve_generative_process_config`, and then resolves again with the same + arguments. Harmless today (the second pass hits the equality branch) but redundant, and it is + the natural thing to unify while restructuring this function rather than adding a fourth + variation of the same block. +- `typed_instantiate` (`utils/config_utils.py:145`) enforces its contract with a bare `assert`, + which is stripped under `python -O`. If it becomes the Protocol check, it should raise. + +### 5.5 What was verified vs. what is still a claim + +Verified by execution against the real code (§3.2, and a Protocol discrimination check): + +- vendored `generators` process runs through `generate_data_batch` correctly; +- exactly two ABC members are missing, both log-space, neither on the training path; +- the current failure is a silent `None` with no warning; +- a `runtime_checkable` core Protocol accepts the vendored adapter **and** the native + `build_hidden_markov_model` output, and rejects a non-conforming object — i.e. the relaxation is + back-compatible by construction, not by hope. + +Not yet verified — the honest gaps in this proposal: + +- a **factored / nonergodic** vendored process end-to-end (the spike used ghmm, the simplest case; + the `encode`/`decode` assembly of §3.4 is where this gets interesting, and `probability`'s + signature genuinely diverges there); +- the full `managed_run` path with MLflow + reproducibility asserts under `strict=True` (the spike + drove `_setup_generative_processes` directly); +- `activation_tracker` against vendored-process beliefs, which is where the MSP gap (§4.3) would + actually bite. + +--- + +## 6 Deprecation path + +`simplexity` is a shared repo with real users and a standing reproducibility obligation — +`generators`' own USAGE.md names it ("because past research must remain reproducible, even internal +APIs become effectively frozen"). The greenfield "no back-compat shims" rule does not apply. +Nothing in `generative_processes` gets deleted or broken by this work. + +The repo's existing deprecation idiom is a comment plus continued function — `pyproject.toml` marks +`aws = [...] # Deprecated S3 Persister.` and `penzai = [...] # Deprecated: penzai is no longer +maintained.` while both keep working. There is **no** `DeprecationWarning` anywhere in +`simplexity/` today, so introducing runtime warnings is a new (and teammate-visible) convention — +hence Q4. + +Proposed staging, each stage independently landable: + +| Stage | Content | Breaks anything? | +| --- | --- | --- | +| **0. Enable** (this PR) | protocol + discovery + loud validation + vendored end-to-end test + migration guide | No. Existing configs and processes untouched; all existing tests are the back-compat proof. | +| **1. Signpost** | `docs/` migration guide, README pointer to `generators`, module docstring notes. Package still fully supported. | No | +| **2. Discourage** | new processes go in consumers, not here. Optionally a `DeprecationWarning` on import of `generative_processes.builder` — **only** once teammates have agreed (Q4), since it fires in everyone's runs. | No, but noisy | +| **3. Re-home the gaps** | §4.3 items 1–3 move out of `generative_processes` into consumer-side homes that operate on the Protocol (`analysis/`, `run_management/`); item 4 resolved per Q5 | Import paths move — needs coordination | +| **4. Thin** | remove what has no consumers, once `eden-experiments` and teammates' branches have migrated | Yes — needs a deprecation window and sign-off | + +Stages 3–4 are explicitly **not** in this PR's scope; naming them is what keeps stage 0 honest +about being a beginning rather than a fait accompli. + +### 6.1 Migration guide sketch (to be written as `docs/`, aimed at a consumer) + +1. `uv add` nothing. Copy the modules you need from `generators` into your project + (`ghmm/process.py` + `utils.py` is the minimal unit; add `factored/*.py` for composites), fix + the module-level imports to your package paths, delete what you do not use. +2. Copy the transition-matrix constructors you need (`transition_matrices/classical.py`), or write + your own — these are data. +3. Write one adapter module binding the vendored functions to the protocol's member names. ~50 + lines for a ghmm; the worked example ships in `simplexity`'s docs and its test fixture. +4. Write a factory in the same module for the config to target (this is where `encode`/`decode` + assembly lives for composite processes). +5. Point your Hydra config at the factory; declare the section as a generative process. +6. Wire the `generators` conformance vectors into your test suite against your copy — **pending + Q3**, since they do not exist yet. +7. Run. `vocab_size` / `bos_token` resolution and the training path work as before. + +A guide that does not cover the actual consumers is not a guide — §7 is the list it must cover. + +--- + +## 7 Consumer inventory + +**In `simplexity` (all keep working; they are the regression proof):** + +- 15 test modules under `tests/generative_processes/`, plus + `tests/structured_configs/test_generative_process_config.py` (24 references — the densest + coupling to the config machinery, and the file most affected by an added config field), + `tests/run_management/test_components.py`, `tests/end_to_end/training.py`. +- 11 process configs in `tests/end_to_end/configs/generative_process/`. +- `simplexity` internals: `run_management/run_management.py`, `run_management/components.py`, + `structured_configs/generative_process.py`. +- `walkthroughs/` — greps clean for `generative_processes`; no migration burden. + +**Outside `simplexity` (`eden-experiments`) — four consumers, and the shape of their usage is the +good news:** + +| Consumer | Imports | Migration cost | +| --- | --- | --- | +| `belief-state-recovery/ambitious/{harness,pipeline_smoke}.py` | `generative_processes.torch_generator.generate_data_batch` | low — a utility that takes a process, not a process. Works on vendored objects today (§3.2 line 4). | +| `fraxl-run/salvage-candidate/{train,analyze}.py` | same | low, same reason | +| `fraxl-run/salvage-candidate/{process.py, scratch/sweep*.py}` | `hidden_markov_model.HiddenMarkovModel`, `mixed_state_presentation.MixedStateTreeGenerator` | **medium** — the MSP gap (§4.3 item 1) | +| `nonergodic-demo/nonergodic_entropy.py` | `builder.*`, `mixed_state_presentation.*` | **medium** — same | + +Three of four use `generate_data_batch`, which is process-agnostic and already vendored-compatible. +The two that would actually be blocked are blocked on MSP, which reinforces that §4.3 item 1 is the +critical-path gap rather than a footnote. Also note `r2-analysis/95_log_to_mlflow.py` uses +`managed_run` **without** a generative-process config, which is the case the current silent skip is +indistinguishable from. + +**Not surveyed:** `origin` carries 60+ branches, several with process work in flight +(`adam/factor_tree`, `adam/hidden-factors`, `adam/producted-generator*`, `casper/*`). Any config +schema change collides with those on merge. This is a coordination cost, not a technical one, and +it is the main reason B3 (additive) is recommended over B2 (migrate every config). + +--- + +## 8 Decisions taken + +Recorded 2026-07-30. §9 lists what is still open. + +**D1 — Reading A (§4.1).** `simplexity` becomes a runner; processes live in consumers' projects. +This document's recommendation, adopted. + +**D2 — Full deprecation is gated on generators reaching feature parity.** The first change ships +a `PendingDeprecationWarning` on `simplexity.generative_processes` immediately, chosen over both +docs-only signposting and a full `DeprecationWarning`. `PendingDeprecationWarning` is ignored by +default in CPython, so it does not add noise to teammates' runs, while still surfacing in pytest +and for anyone running with warnings enabled. + +The gating clause has a consequence worth making explicit, because it revises §4.3 of this +document: if full deprecation waits on *generators* achieving parity, then the four capability gaps +are **upstream generators work**, not simplexity re-homing. §4.3 originally proposed moving +`mixed_state_presentation.py` into a stable simplexity home; under D2 the primary path is instead +that generators grows the capability. See Q1 in §9 — the mixed-state tree is the case where this +distinction genuinely bites, because it is arguably belief-state *analysis* rather than a process +operation, and analysis is not obviously in generators' scope. + +**D3 — One PR to `main`**, carrying design and implementation together, held to `main`'s bar: new +tests for new behaviour, coverage of failure modes, all static checks passing, no suppressions +without justification. + +**D4 — Protocol named `GenerativeProcessProtocol`.** I had shipped it as `GenerativeProcess`, +reusing the name so the call site would read correctly once the ABC retires. Overruled, correctly: +that argument optimizes an end state gated on Q1, which is unresolved and is the largest open item +here — so the transition window is of unbounded length, and throughout it two classes named +`GenerativeProcess` live in two modules. My own defence, that the ambiguity is "confined to readers +rather than to code", is the argument *for* renaming: readers are who naming serves, and the readers +are teammates on 60+ branches reading tracebacks, error strings and review diffs where the bare name +does not say which of the two it is. The usual counter — that Python's protocols carry no `Protocol` +suffix (`Iterable`, `Sequence`) — does not apply, because none of those coexist with a same-named ABC +in the same codebase. The collision decides it, not the convention. Cost is asymmetric too: the +protocol is referenced from one module, so a later rename is trivial, while the collision is paid +continuously. + +The extension protocol was renamed to `LogSpaceGenerativeProcessProtocol` for the same reason applied +one step further: it does not itself collide with anything, but a suffixed and an unsuffixed name +sitting side by side in one module tells a reader nothing about which is a protocol. + +**D5 — The declaration field stays `component`.** It matches the runner's existing vocabulary +(`Components`, `component_name`); `component_type` is genuinely taken by factored-process specs +(`hmm` / `ghmm`); `kind` is vaguer; and `external: true` names the exception rather than the role, so +it ages badly exactly when declaration becomes the norm. + +## 9 Open questions + +Renumbered as questions are resolved; D1-D5 above are settled. + +**Q1 — Does the mixed-state tree belong in generators or in simplexity?** Under D2 the gaps close +upstream, but `mixed_state_presentation.py` (576 lines) enumerates the belief-state tree and +computes myopic entropy: that is analysis *over* a process, not an operation *of* one, and +generators' SPEC covers belief updates rather than tree enumeration. Two `eden-experiments` +consumers import it directly. If it is not going upstream, parity is unreachable by definition and +it needs a stable simplexity home instead (`simplexity/analysis/` operating on the protocol). +This is the single largest open item and it gates the deprecation timeline under D2. + +**Q2 — Conformance vectors still do not exist.** SPEC §7 specifies the encoding; no fixtures ship +in generators. Under a structural contract they are the only thing that keeps a consumer's modified +copy mathematically honest, because the protocol check verifies that members exist and nothing +about what they compute (§5.1). The migration guide currently tells consumers to wire up vectors +and then admits they are unavailable. Is generating them upstream work you want to schedule, or +should simplexity offer a conformance harness against its own implementations as an interim? + +**Q3 — Does `transition_matrices.py` get deprecated too?** 16 process families, and it is *data* +rather than behaviour — the most reusable and least contentious part of the package. `generators` +covers ~5 and its USAGE.md disclaims owning a process zoo while shipping `transition_matrices/` +anyway. Options: keep it in `simplexity` as data (my lean), upstream the missing families, or have +each consumer vendor what it uses. Under D2 this also affects what "parity" means. + +**Q4 — Should existing configs adopt the declaration?** Discovery keeps the namespace prefix as a +fast path, so the 11 in-repo configs and teammates' in-flight branches needed no changes. Under +Reading A every process config eventually declares itself and the prefix path retires. Migrating +the in-repo configs now is a small mechanical change that would make the intent uniform; leaving +them is less churn against the 60+ open branches. Deferred to you. + +**Q5 — Where does `generators` live?** The repository is currently `ealt/generators` (public); +`Astera-org/generators` does not exist. Every reference in this change points at `ealt/generators` +because that is what resolves today. If the intent is for it to become an Astera-org repository, +that should happen before consumers start vendoring from it in earnest: a shared Astera repo whose +deprecation notice points at an individual's account is awkward, and the URL appears in a runtime +warning message that teammates will see. + +**Q6 — Two lint errors and 20 pyright errors pre-date this branch.** +`tests/generative_processes/test_data_prefetcher.py` has 3 ruff findings (SIM117 ×2, PT012) at +`HEAD`, and pyright reports 20 unresolved-import errors for the uninstalled `penzai` and `aws` +optional extras. Both are in files this change does not touch, so they are left alone rather than +folded into an unrelated diff. `main`'s bar says all checks must pass, so someone should decide +whether that is a separate cleanup PR or an accepted baseline. + +## 10 Appendix: what the implementation changed + +| File | Change | +| --- | --- | +| `simplexity/run_management/protocols.py` | New. `GenerativeProcessProtocol` and `LogSpaceGenerativeProcessProtocol`, plus `missing_generative_process_members` for actionable errors. | +| `simplexity/utils/config_utils.py` | `filter_instance_keys_by` takes a config-aware predicate; `filter_instance_keys` reimplemented over it, unchanged for callers. Added `get_instance_target`. | +| `simplexity/structured_configs/generative_process.py` | `component` field; `declares_generative_process`, `declares_generative_process_instance_key`, `claims_generative_process_instance_key`; validation accepts declared foreign targets and explains the declaration when rejecting. | +| `simplexity/run_management/run_management.py` | Discovery via the claims predicate (both in setup and vocabulary resolution); protocol check replacing the nominal `typed_instantiate`; declared-but-invalid now raises; the not-found log names what it considered; removed the duplicated resolve. | +| `simplexity/run_management/components.py` | `generative_processes` typed by the protocol. | +| `simplexity/generative_processes/generator.py`, `torch_generator.py` | Accept any conforming process. Narrowed `tokens` at the `generate` boundary, which also cleared pre-existing pyright looseness that `filter_vmap` had been masking. | +| `simplexity/generative_processes/__init__.py` | New. `PendingDeprecationWarning`. | +| `docs/generators_migration.md` | New. Consumer-facing migration guide. | +| `tests/vendored_process/` | New. Verbatim generators copies at `b1242ea`, plus the adapter that is the worked example. | +| `tests/end_to_end/test_vendored_process_training.py` | New. The acceptance test: a vendored process trains under `@managed_run()`. | +| `pyproject.toml` | `pythonpath = ["."]` so tests can import the fixture; ruff ignores docstring rules for the verbatim vendored copies. | + +One friction point worth recording, because every consumer will hit it: **vendored code arrives with +the upstream repo's conventions.** Generators' `ZEN.md` says "Eschew excess documentation," so its +modules have no docstrings and fail simplexity's `D` lint rules. Editing them to comply would +defeat the ability to diff against upstream, so the vendored directory is excluded from those rules +instead. Consumers vendoring generators into their own projects will need the same exclusion. diff --git a/docs/generators_migration.md b/docs/generators_migration.md new file mode 100644 index 00000000..94819ea0 --- /dev/null +++ b/docs/generators_migration.md @@ -0,0 +1,149 @@ +# Migrating to a vendored generative process + +Generative processes are moving out of simplexity and into your own project, sourced from +[generators](https://github.com/ealt/generators). This guide shows how to run a vendored +process under `@managed_run()`. + +Nothing is removed yet. `simplexity.generative_processes` remains fully supported and emits only a +`PendingDeprecationWarning`; full deprecation waits until generators reaches feature parity. If your +code works today, it keeps working, and you can migrate when it suits you. + +## Why copy instead of import + +Generators is not a library and should not be added to your dependencies. You copy the modules you +need into your project, delete what you do not use, and modify the rest until it fits your problem. +The process code then lives where you can read, instrument, and optimize it, and your experiment +artifacts carry the exact code that produced them. The invariant that survives your edits is not +the code but the specification: generators' `SPEC.md` defines the required results, and its +conformance vectors are how you check that your copy still implements the same mathematics. + +Simplexity supports this by identifying a generative process by the operations it provides rather +than by where its code lives. + +## The contract + +Your process must provide `simplexity.run_management.protocols.GenerativeProcessProtocol`: + +| Member | Meaning | +| --- | --- | +| `vocab_size` | Number of emittable observations. Used to resolve `vocab_size` in your config and `d_vocab` in your model's. | +| `initial_state` | The state generation starts from. | +| `emit_observation(state, key)` | Sample one observation. | +| `transition_states(state, obs)` | Update the state after an observation. | +| `observation_probability_distribution(state)` | Distribution over observations from a state. | +| `probability(observations)` | Probability of a sequence. | +| `generate(state, key, sequence_len, return_all_states)` | Generate a batch of sequences, optionally returning every belief state. | + +It is a `Protocol`, so you do not inherit from anything — an object with these members conforms. + +Belief-state analysis additionally needs `LogSpaceGenerativeProcessProtocol` +(`log_observation_probability_distribution`, `log_probability`). Generators does not provide +log-space operations, so add them only if your run does belief analysis; training and generation +never call them. + +## Steps + +### 1. Copy the modules + +Copy what you need from generators into your project, fix the module-level imports to your own +paths, and delete the rest. `ghmm/process.py` plus `utils.py` is the minimal unit; add +`factored/*.py` for composite processes. Copy the transition-matrix constructors you want from +`transition_matrices/`, or write your own — those are data, not behaviour. + +Record the upstream commit you copied from, so you can diff against it later to pick up +improvements. + +### 2. Write an adapter + +Generators exposes module-level functions over a `NamedTuple` of process data, so bind that data +once and expose the operations as methods. `tests/vendored_process/adapter.py` in this repo is a +complete worked example, around fifty lines for a GHMM. + +Two things usually need attention: + +- **Generation.** Generators' `generate` returns only the final state. Simplexity's training path + also wants every intermediate belief state, so re-scan rather than calling it, and batch with + `eqx.filter_vmap` (or an explicit `jax.vmap`) as the example does. +- **Composite processes.** `factored/*.py` takes keyword-only `encode` and `decode` callables + because generators deliberately does not own the composite-token encoding. Bind your choice in + the adapter. Note also that `factored`'s `seq_prob` takes a state, unlike the GHMM's, so your + `probability` should bind the initial state. + +### 3. Write a factory + +Assembling a process is code — constructing matrices, calling `init`, binding `encode`/`decode` — +and none of that fits in YAML. Give your adapter module a factory function for the config to +target: + +```python +def build_my_process(x: float, a: float) -> MyVendoredProcess: + return MyVendoredProcess(data=ghmm_process.init(mess(x, a, 3))) +``` + +This mirrors how simplexity's own `builder.build_hidden_markov_model` is configured. + +### 4. Declare it in your config + +Point `_target_` at your factory, and set `component: generative_process` so run management knows +what the section configures. The declaration is required for processes outside +`simplexity.generative_processes`: your process lives in your own namespace, so it cannot be +recognized from its import path. + +```yaml +name: my_vendored_process +component: generative_process +instance: + _target_: my_project.processes.adapter.build_my_process + x: 0.15 + a: 0.6 + +base_vocab_size: ??? +bos_token: ??? +eos_token: null +vocab_size: ??? +``` + +The `???` fields resolve from your process's `vocab_size`, exactly as for a simplexity process, and +feed `d_vocab` in your model's config. + +### 5. Wire up conformance vectors + +Copy the relevant conformance vectors from generators into your test suite and run them against +your copy: green means your fork still implements the specification. + +Be aware that generators specifies the vector encoding in `SPEC.md` section 7 but does not yet ship +the vectors themselves. Until it does, this safety net is unavailable, and the protocol check will +not cover it — `isinstance` against a protocol verifies that members exist, not that they compute +the right thing. Until vectors exist, test your copy against whatever ground truth you have. + +### 6. Run + +Your entrypoint is unchanged. `components.get_generative_process()` returns your process, and the +generation helpers accept it: + +```python +from simplexity.generative_processes.torch_generator import generate_data_batch +``` + +## Troubleshooting + +**"no generative process configs found", and `get_generative_process()` returns `None`.** Your +section was not claimed. Check that `component: generative_process` is set on the section itself, +not on the nested `instance`. The log message lists the instance keys that were considered. + +**"config declares generative processes at ..., but their configs are invalid".** The section +declared itself but failed validation. The preceding validation warnings say why. This is +deliberately an error rather than a skip: a declared process that quietly does not load used to +leave the run with no process at all. + +**"... is not a generative process: missing ...".** Your factory returned something that does not +provide the whole contract; the message names the absent members. A common cause is a typo in a +method name, since conformance is structural. + +## Reference + +- `docs/design/generators_instantiation.md` — the design, the tradeoffs considered, and the + outstanding gaps between generators and this package. +- `tests/vendored_process/` — a checked-in vendored process and its adapter. +- `tests/end_to_end/test_vendored_process_training.py` — that process training under + `@managed_run()`. diff --git a/pyproject.toml b/pyproject.toml index 3f47c89c..2394e610 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,6 +92,10 @@ convention = "google" [tool.ruff.lint.per-file-ignores] "test_*.py" = ["D"] "*.ipynb" = ["D"] +# Verbatim copies of upstream generators modules, which follow that repo's conventions +# (ZEN.md: "Eschew excess documentation"). Keeping them unedited is what makes diffing +# against upstream to pick up improvements possible. +"tests/vendored_process/generators_copy/*" = ["D"] [tool.pylint.MASTER] # Add project root to Python path so pylint can resolve simplexity imports @@ -101,7 +105,19 @@ load-plugins = ["pylint_per_file_ignores"] [tool.pylint."MESSAGES CONTROL"] # Disable line-too-long check since ruff handles line length (120 chars) disable = ["line-too-long", "duplicate-code"] -per-file-ignores = ["tests/**:redefined-outer-name"] +per-file-ignores = [ + "tests/**:redefined-outer-name", + # Verbatim copies of upstream generators modules. They follow that repo's conventions -- + # no docstrings (ZEN.md: "Eschew excess documentation") and notation-faithful names that + # mirror the mathematics (VARIABLE_NAMING.md: Ts, T, Vs). Editing them to satisfy this + # repo's style would defeat diffing against upstream to pick up improvements. + "tests/vendored_process/generators_copy/**:missing-module-docstring,missing-class-docstring,missing-function-docstring,invalid-name", + # Pre-existing debt, newly in scope: this subpackage had no __init__.py, so `pylint simplexity` + # never walked into it. Adding one (to carry the pending-deprecation warning) brought 21 modules + # into scope at once. The package is pending deprecation, so refactoring it is not worthwhile; + # tracked in https://github.com/Astera-org/simplexity/issues/199. + "simplexity/generative_processes/**:too-many-arguments,too-many-locals,too-many-instance-attributes,unused-argument,invalid-name", +] [tool.pylint.IMPORTS] # Ignore simplexity module to prevent astroid crashes during import resolution @@ -130,6 +146,8 @@ addopts = [ # To check overall coverage locally, use: pytest --cov-fail-under=80 ] markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"] +# Makes the repo root importable so tests can import vendored process fixtures under tests/. +pythonpath = ["."] [tool.coverage.run] source = ["simplexity"] @@ -156,6 +174,8 @@ exclude_lines = [ "if __name__ == .__main__.:", "if TYPE_CHECKING:", "@abstractmethod", + # Protocol method bodies, the structural counterpart to @abstractmethod above. + '^\s*\.\.\.$', ] show_missing = true skip_covered = false diff --git a/simplexity/generative_processes/__init__.py b/simplexity/generative_processes/__init__.py new file mode 100644 index 00000000..8541e88c --- /dev/null +++ b/simplexity/generative_processes/__init__.py @@ -0,0 +1,21 @@ +"""Generative process implementations, pending deprecation. + +Generative processes are moving out of simplexity and into consumers' own projects, sourced from +[generators](https://github.com/ealt/generators), which distributes process modules by +copying rather than by import. Run management instantiates a vendored process on equal terms with +one implemented here: see `simplexity.run_management.protocols.GenerativeProcessProtocol` for the +and contract and `docs/generators_migration.md` for the migration. + +Nothing here is removed or altered yet. Full deprecation waits until generators reaches feature +parity with this package; `docs/design/generators_instantiation.md` tracks the outstanding gaps. +""" + +import warnings + +warnings.warn( + "simplexity.generative_processes is pending deprecation in favour of processes vendored from " + "generators (https://github.com/ealt/generators). It remains fully supported until " + "generators reaches feature parity; see docs/generators_migration.md.", + PendingDeprecationWarning, + stacklevel=2, +) diff --git a/simplexity/generative_processes/generator.py b/simplexity/generative_processes/generator.py index 32dcd766..53267c05 100644 --- a/simplexity/generative_processes/generator.py +++ b/simplexity/generative_processes/generator.py @@ -15,14 +15,14 @@ import jax import jax.numpy as jnp -from simplexity.generative_processes.generative_process import GenerativeProcess from simplexity.generative_processes.nonergodic_generative_process import NonErgodicState +from simplexity.run_management.protocols import GenerativeProcessProtocol @eqx.filter_jit def generate_data_batch( gen_states: jax.Array | tuple[jax.Array, ...], - data_generator: GenerativeProcess, + data_generator: GenerativeProcessProtocol, batch_size: int, sequence_len: int, key: jax.Array, @@ -31,7 +31,8 @@ def generate_data_batch( ) -> tuple[jax.Array | tuple[jax.Array, ...], jax.Array, jax.Array]: """Generate a batch of data without tracking intermediate beliefs.""" batch_keys = jax.random.split(key, batch_size) - gen_states, tokens = data_generator.generate(gen_states, batch_keys, sequence_len, False) + gen_states, generated = data_generator.generate(gen_states, batch_keys, sequence_len, False) + tokens = jnp.asarray(generated) if bos_token is not None: tokens = jnp.concatenate([jnp.full((batch_size, 1), bos_token), tokens], axis=1) @@ -46,7 +47,7 @@ def generate_data_batch( @eqx.filter_jit def generate_data_batch_with_full_history( gen_states: jax.Array | tuple[jax.Array, ...], - data_generator: GenerativeProcess, + data_generator: GenerativeProcessProtocol, batch_size: int, sequence_len: int, key: jax.Array, @@ -55,7 +56,8 @@ def generate_data_batch_with_full_history( ) -> dict[str, jax.Array | tuple[jax.Array, ...]]: """Generate sequences plus per-token belief states and prefix probabilities.""" batch_keys = jax.random.split(key, batch_size) - belief_states, tokens = data_generator.generate(gen_states, batch_keys, sequence_len, True) + belief_states, generated = data_generator.generate(gen_states, batch_keys, sequence_len, True) + tokens = jnp.asarray(generated) prefix_probs = _compute_prefix_probabilities(data_generator, gen_states, tokens) @@ -119,7 +121,7 @@ def _slice_belief_states( def _compute_prefix_probabilities( - data_generator: GenerativeProcess, + data_generator: GenerativeProcessProtocol, initial_states: jax.Array | tuple[jax.Array, ...], tokens: jax.Array, ) -> jax.Array: diff --git a/simplexity/generative_processes/torch_generator.py b/simplexity/generative_processes/torch_generator.py index 00f4211a..0bbf6672 100644 --- a/simplexity/generative_processes/torch_generator.py +++ b/simplexity/generative_processes/torch_generator.py @@ -12,19 +12,19 @@ import jax import torch -from simplexity.generative_processes.generative_process import GenerativeProcess from simplexity.generative_processes.generator import ( generate_data_batch as generate_jax_data_batch, ) from simplexity.generative_processes.generator import ( generate_data_batch_with_full_history as generate_jax_data_batch_with_full_history, ) +from simplexity.run_management.protocols import GenerativeProcessProtocol from simplexity.utils.pytorch_utils import jax_to_torch def generate_data_batch( gen_states: jax.Array | tuple[jax.Array, ...], - data_generator: GenerativeProcess, + data_generator: GenerativeProcessProtocol, batch_size: int, sequence_len: int, key: jax.Array, @@ -61,7 +61,7 @@ def generate_data_batch( def generate_data_batch_with_full_history( gen_states: jax.Array | tuple[jax.Array, ...], - data_generator: GenerativeProcess, + data_generator: GenerativeProcessProtocol, batch_size: int, sequence_len: int, key: jax.Array, diff --git a/simplexity/run_management/components.py b/simplexity/run_management/components.py index 1b32c169..91350550 100644 --- a/simplexity/run_management/components.py +++ b/simplexity/run_management/components.py @@ -13,10 +13,10 @@ from typing import Any from simplexity.activations.activation_tracker import ActivationTracker -from simplexity.generative_processes.generative_process import GenerativeProcess from simplexity.logging.logger import Logger from simplexity.metrics.metric_tracker import MetricTracker from simplexity.persistence.model_persister import ModelPersister +from simplexity.run_management.protocols import GenerativeProcessProtocol @dataclass @@ -24,7 +24,7 @@ class Components: """Components for the run.""" loggers: dict[str, Logger] | None = None - generative_processes: dict[str, GenerativeProcess] | None = None + generative_processes: dict[str, GenerativeProcessProtocol] | None = None persisters: dict[str, ModelPersister] | None = None predictive_models: dict[str, Any] | None = None # TODO: improve typing optimizers: dict[str, Any] | None = None # TODO: improve typing @@ -36,7 +36,7 @@ def get_logger(self, key: str | None = None) -> Logger | None: """Get the logger.""" return self._get_instance_by_key(self.loggers, key, "logger") - def get_generative_process(self, key: str | None = None) -> GenerativeProcess | None: + def get_generative_process(self, key: str | None = None) -> GenerativeProcessProtocol | None: """Get the generative process.""" return self._get_instance_by_key(self.generative_processes, key, "generative process") diff --git a/simplexity/run_management/protocols.py b/simplexity/run_management/protocols.py new file mode 100644 index 00000000..8f7cb196 --- /dev/null +++ b/simplexity/run_management/protocols.py @@ -0,0 +1,116 @@ +"""Structural contracts for components that run management instantiates. + +A component is identified by what it can do, not by where its code lives. This matters for +generative processes: the reference implementations in +[generators](https://github.com/ealt/generators) are distributed by copying rather than +by import, so a conforming process may be a project-local module that shares no namespace or +base class with simplexity. + +The operations mirror the generators specification, which fixes required results rather than +their organization (SPEC.md sections 2 and 6.1). `GenerativeProcessProtocol` covers the operations the +runner and the training path exercise; log-space operations are a separate protocol because +only belief-state analysis needs them. +""" + +from typing import Any, Protocol, runtime_checkable + +import chex +import jax + + +def _protocol_member_names(protocol: type) -> frozenset[str]: + """The public member names a protocol declares, including inherited ones. + + Walks the MRO rather than reading the protocol's `__protocol_attrs__`, which is a CPython + implementation detail that static analysis does not model. + """ + return frozenset( + name + for klass in protocol.__mro__ + if klass is not object and klass.__module__ != "typing" + for name in vars(klass) + if not name.startswith("_") + ) + + +@runtime_checkable +class GenerativeProcessProtocol(Protocol): + """A probabilistic model over observation sequences that run management can drive. + + Named for the contract rather than the concept, because the concept's other name is taken: + `simplexity.generative_processes.generative_process.GenerativeProcess` is the abstract base + class that in-repo processes inherit from. This is the structural contract they satisfy, which + a process implemented anywhere can satisfy equally. + + Implementations may be subclasses of that base class or project-local modules vendored from + generators and adapted. Only the members below are required. + """ + + @property + def vocab_size(self) -> int: + """The number of observations that can be emitted by the generative process.""" + ... + + @property + def initial_state(self) -> Any: + """The initial state of the generative process.""" + ... + + def emit_observation(self, state: Any, key: chex.PRNGKey) -> chex.Array: + """Emit an observation based on the state of the generative process.""" + ... + + def transition_states(self, state: Any, obs: chex.Array) -> Any: + """Evolve the state of the generative process based on the observation.""" + ... + + def observation_probability_distribution(self, state: Any) -> jax.Array: + """Compute the distribution over observations that can be emitted from a state.""" + ... + + def probability(self, observations: jax.Array) -> jax.Array: + """Compute the probability of the process generating a sequence of observations.""" + ... + + def generate( + self, state: Any, key: chex.PRNGKey, sequence_len: int, return_all_states: bool + ) -> tuple[Any, chex.Array]: + """Generate a batch of observation sequences, optionally returning every belief state.""" + ... + + +@runtime_checkable +class LogSpaceGenerativeProcessProtocol(GenerativeProcessProtocol, Protocol): + """A generative process that also exposes log-space operations. + + Required only by belief-state analysis, such as mixed-state presentation, where working in + log space keeps long-sequence probabilities numerically stable. The training and generation + paths never call these. + """ + + def log_observation_probability_distribution(self, log_belief_state: Any) -> jax.Array: + """Compute the log distribution over observations that can be emitted from a state.""" + ... + + def log_probability(self, observations: jax.Array) -> jax.Array: + """Compute the log probability of the process generating a sequence of observations.""" + ... + + +GENERATIVE_PROCESS_MEMBERS = _protocol_member_names(GenerativeProcessProtocol) +LOG_SPACE_GENERATIVE_PROCESS_MEMBERS = _protocol_member_names(LogSpaceGenerativeProcessProtocol) + + +def missing_generative_process_members(obj: object) -> list[str]: + """List the `GenerativeProcessProtocol` members that an object does not provide. + + `isinstance` against a runtime-checkable protocol reports only whether an object conforms. + Reporting *which* members are absent turns a rejected config into an actionable error. + + Args: + obj: The candidate generative process. + + Returns: + The names of missing members, sorted. + """ + return sorted(name for name in GENERATIVE_PROCESS_MEMBERS if not hasattr(obj, name)) diff --git a/simplexity/run_management/run_management.py b/simplexity/run_management/run_management.py index bcb046a7..a4e86e80 100644 --- a/simplexity/run_management/run_management.py +++ b/simplexity/run_management/run_management.py @@ -36,13 +36,14 @@ from omegaconf import DictConfig, OmegaConf from torch.nn import Module as PytorchModel -from simplexity.generative_processes.generative_process import GenerativeProcess +from simplexity.exceptions import ConfigValidationError from simplexity.logger import SIMPLEXITY_LOGGER, add_handlers_to_existing_loggers, get_log_files, remove_log_files from simplexity.logging.logger import Logger from simplexity.logging.mlflow_logger import MLFlowLogger from simplexity.persistence.mlflow_persister import MLFlowPersister from simplexity.persistence.model_persister import ModelPersister from simplexity.run_management.components import Components +from simplexity.run_management.protocols import GenerativeProcessProtocol, missing_generative_process_members from simplexity.run_management.run_logging import ( log_environment_artifacts, log_git_info, @@ -56,7 +57,9 @@ ) from simplexity.structured_configs.base import resolve_base_config, validate_base_config from simplexity.structured_configs.generative_process import ( - is_generative_process_target, + GENERATIVE_PROCESS_COMPONENT, + claims_generative_process_instance_key, + declares_generative_process_instance_key, resolve_generative_process_config, validate_generative_process_config, ) @@ -91,8 +94,10 @@ ) from simplexity.utils.config_utils import ( filter_instance_keys, + filter_instance_keys_by, get_config, get_instance_keys, + get_instance_target, typed_instantiate, ) from simplexity.utils.jnp_utils import resolve_jax_device @@ -327,45 +332,76 @@ def _setup_logging(cfg: DictConfig, instance_keys: list[str], *, strict: bool) - return None -def _instantiate_generative_process(cfg: DictConfig, instance_key: str) -> GenerativeProcess: - """Setup the generative process.""" +def _instantiate_generative_process(cfg: DictConfig, instance_key: str) -> GenerativeProcessProtocol: + """Instantiate a generative process and resolve its config's vocabulary fields. + + The process is accepted on the strength of the operations it provides rather than the class it + inherits from, so that a process vendored into the consumer's project conforms on equal terms + with one implemented in simplexity. + """ instance_config = OmegaConf.select(cfg, instance_key, throw_on_missing=True) - if instance_config: - generative_process = typed_instantiate(instance_config, GenerativeProcess) - SIMPLEXITY_LOGGER.info( - "[generative process] instantiated generative process: %s", generative_process.__class__.__name__ + if instance_config is None: + raise KeyError(instance_key) + generative_process = hydra.utils.instantiate(instance_config) + missing_members = missing_generative_process_members(generative_process) + if missing_members: + raise ConfigValidationError( + f"{instance_key} instantiated {type(generative_process).__name__}, which is not a generative process: " + f"missing {', '.join(missing_members)}" ) - config_key = instance_key.rsplit(".", 1)[0] - generative_process_config: DictConfig | None = OmegaConf.select(cfg, config_key) - if generative_process_config is None: - raise RuntimeError("Error selecting generative process config") - base_vocab_size = generative_process.vocab_size - resolve_generative_process_config(generative_process_config, base_vocab_size) - return generative_process - raise KeyError + SIMPLEXITY_LOGGER.info( + "[generative process] instantiated generative process: %s", generative_process.__class__.__name__ + ) + config_key = instance_key.rsplit(".", 1)[0] + generative_process_config: DictConfig | None = OmegaConf.select(cfg, config_key) + if generative_process_config is None: + raise RuntimeError("Error selecting generative process config") + resolve_generative_process_config(generative_process_config, generative_process.vocab_size) + return generative_process -def _setup_generative_processes(cfg: DictConfig, instance_keys: list[str]) -> dict[str, GenerativeProcess] | None: - instance_keys = filter_instance_keys( +def _assert_declared_processes_were_claimed(cfg: DictConfig, instance_keys: list[str], claimed: list[str]) -> None: + """Fail loudly when a section declares a generative process that discovery then discarded. + + Discovery drops config sections it cannot validate, which is the right default for an + unrecognized target but wrong for a section that explicitly declared its intent: silently + proceeding leaves `components.generative_processes` as None, indistinguishable from a config + that configures no process at all. + """ + discarded = [ + instance_key + for instance_key in instance_keys + if instance_key not in claimed and declares_generative_process_instance_key(cfg, instance_key) + ] + if discarded: + raise ConfigValidationError( + f"config declares generative processes at {', '.join(discarded)}, but their configs are invalid; " + "see the preceding validation warnings" + ) + + +def _setup_generative_processes( + cfg: DictConfig, instance_keys: list[str] +) -> dict[str, GenerativeProcessProtocol] | None: + claimed_instance_keys = filter_instance_keys_by( cfg, instance_keys, - is_generative_process_target, + claims_generative_process_instance_key, validate_fn=validate_generative_process_config, component_name="generative process", ) - if instance_keys: - generative_processes = {} - for instance_key in instance_keys: - generative_process = _instantiate_generative_process(cfg, instance_key) - config_key = instance_key.rsplit(".", 1)[0] - generative_process_config: DictConfig | None = OmegaConf.select(cfg, config_key) - if generative_process_config is None: - raise RuntimeError("Error selecting generative process config") - base_vocab_size = generative_process.vocab_size - resolve_generative_process_config(generative_process_config, base_vocab_size) - generative_processes[instance_key] = generative_process - return generative_processes - SIMPLEXITY_LOGGER.info("[generative process] no generative process configs found") + _assert_declared_processes_were_claimed(cfg, instance_keys, claimed_instance_keys) + if claimed_instance_keys: + return { + instance_key: _instantiate_generative_process(cfg, instance_key) for instance_key in claimed_instance_keys + } + unclaimed = {instance_key: get_instance_target(cfg, instance_key) for instance_key in instance_keys} + SIMPLEXITY_LOGGER.info( + "[generative process] no generative process configs found; considered %s. A process implemented outside " + "simplexity must set component: %s on its config section to be recognized.", + unclaimed, + GENERATIVE_PROCESS_COMPONENT, + ) return None @@ -408,10 +444,10 @@ def _get_persister(persisters: dict[str, ModelPersister] | None) -> ModelPersist def _get_attribute_value(cfg: DictConfig, instance_keys: list[str], attribute_name: str) -> int | None: """Get the vocab size.""" - instance_keys = filter_instance_keys( + instance_keys = filter_instance_keys_by( cfg, instance_keys, - is_generative_process_target, + claims_generative_process_instance_key, validate_fn=validate_generative_process_config, component_name="generative process", ) diff --git a/simplexity/structured_configs/generative_process.py b/simplexity/structured_configs/generative_process.py index 06aaf09e..9706ae44 100644 --- a/simplexity/structured_configs/generative_process.py +++ b/simplexity/structured_configs/generative_process.py @@ -30,7 +30,7 @@ validate_sequence, validate_transition_matrices, ) -from simplexity.utils.config_utils import dynamic_resolve +from simplexity.utils.config_utils import dynamic_resolve, get_instance_target @dataclass @@ -372,12 +372,23 @@ def validate_hidden_markov_model_instance_config(cfg: DictConfig) -> None: validate_initial_state(initial_state, transition_matrices.shape[1], "HiddenMarkovModelInstanceConfig.initial_state") +GENERATIVE_PROCESS_COMPONENT = "generative_process" + + @dataclass class GenerativeProcessConfig: - """Base configuration for generative processes.""" + """Base configuration for generative processes. + + Attributes: + component: Set to `GENERATIVE_PROCESS_COMPONENT` to declare that this section configures a + generative process whose implementation lives outside simplexity, such as a project-local + module vendored from generators. Sections whose `instance._target_` is a simplexity + generative process are recognized without it. + """ instance: InstanceConfig name: str | None = None + component: str | None = None base_vocab_size: int = MISSING bos_token: int | None = MISSING eos_token: int | None = MISSING @@ -398,6 +409,35 @@ def is_generative_process_config(cfg: DictConfig) -> bool: return False +def declares_generative_process(cfg: DictConfig) -> bool: + """Check if a config section declares itself to configure a generative process. + + Args: + cfg: The component's config section, not its nested instance config. + """ + return cfg.get("component", None) == GENERATIVE_PROCESS_COMPONENT + + +def declares_generative_process_instance_key(cfg: DictConfig, instance_key: str) -> bool: + """Check if the section owning an instance key declares itself a generative process.""" + section = OmegaConf.select(cfg, instance_key.rsplit(".", 1)[0], throw_on_missing=False) + return isinstance(section, DictConfig) and declares_generative_process(section) + + +def claims_generative_process_instance_key(cfg: DictConfig, instance_key: str) -> bool: + """Check if an instance key configures a generative process. + + A target under `simplexity.generative_processes` identifies a process by itself. Any other + target must be claimed by an explicit declaration on the owning section, because processes + vendored from generators live in the consumer's own namespace and cannot be recognized from + their import path. + """ + target = get_instance_target(cfg, instance_key) + if target is not None and is_generative_process_target(target): + return True + return declares_generative_process_instance_key(cfg, instance_key) + + def validate_generative_process_config(cfg: DictConfig) -> None: """Validate a GenerativeProcessConfig. @@ -423,8 +463,11 @@ def validate_generative_process_config(cfg: DictConfig) -> None: validate_hidden_markov_model_instance_config(instance) else: validate_instance_config(instance) - if not is_generative_process_config(instance): - raise ConfigValidationError("GenerativeProcessConfig.instance must be a generative process target") + if not is_generative_process_config(instance) and not declares_generative_process(cfg): + raise ConfigValidationError( + "GenerativeProcessConfig.instance must be a generative process target, or the section must set " + f"component: {GENERATIVE_PROCESS_COMPONENT} to declare a process implemented outside simplexity" + ) validate_nonempty_str(name, "GenerativeProcessConfig.name", is_none_allowed=True) _base_vocab_size: int | None = None diff --git a/simplexity/utils/config_utils.py b/simplexity/utils/config_utils.py index 3d96857a..bd19abb8 100644 --- a/simplexity/utils/config_utils.py +++ b/simplexity/utils/config_utils.py @@ -53,6 +53,42 @@ def _validate( return True +def get_instance_target(cfg: DictConfig, instance_key: str) -> str | None: + """Get the `_target_` of an instance key, or None if it is absent or not a string.""" + target = OmegaConf.select(cfg, f"{instance_key}._target_", throw_on_missing=False) + return target if isinstance(target, str) else None + + +def filter_instance_keys_by( + cfg: DictConfig, + instance_keys: list[str], + claims_fn: Callable[[DictConfig, str], bool], + validate_fn: Callable[[DictConfig], None] | None = None, + component_name: str | None = None, +) -> list[str]: + """Filter instance keys by a predicate over the whole config. + + Components whose implementations may live outside simplexity cannot be identified from the + `_target_` string alone, so the predicate receives the config and the instance key and may + consult anything it needs, such as an explicit declaration on the component's config section. + + Args: + cfg: The full config. + instance_keys: Candidate instance keys. + claims_fn: Predicate deciding whether an instance key belongs to this component. + validate_fn: Optional validator applied to the instance key's parent config section. + component_name: Component name used to prefix log messages. + + Returns: + The instance keys claimed by `claims_fn` whose configs validate. + """ + return [ + instance_key + for instance_key in instance_keys + if claims_fn(cfg, instance_key) and _validate(cfg, instance_key, validate_fn, component_name) + ] + + def filter_instance_keys( cfg: DictConfig, instance_keys: list[str], @@ -61,12 +97,12 @@ def filter_instance_keys( component_name: str | None = None, ) -> list[str]: """Filter instance keys by filter function to their targets.""" - filtered_instance_keys: list[str] = [] - for instance_key in instance_keys: - target = OmegaConf.select(cfg, f"{instance_key}._target_", throw_on_missing=False) - if isinstance(target, str) and filter_fn(target) and _validate(cfg, instance_key, validate_fn, component_name): - filtered_instance_keys.append(instance_key) - return filtered_instance_keys + + def claims_fn(cfg: DictConfig, instance_key: str) -> bool: + target = get_instance_target(cfg, instance_key) + return target is not None and filter_fn(target) + + return filter_instance_keys_by(cfg, instance_keys, claims_fn, validate_fn, component_name) def get_config(args: tuple[Any, ...], kwargs: dict[str, Any]) -> DictConfig: diff --git a/tests/end_to_end/configs/generative_process/vendored_mess3.yaml b/tests/end_to_end/configs/generative_process/vendored_mess3.yaml new file mode 100644 index 00000000..af66263e --- /dev/null +++ b/tests/end_to_end/configs/generative_process/vendored_mess3.yaml @@ -0,0 +1,16 @@ +# A generative process vendored from generators into the consumer's own project. +# +# The target is a project-local factory, not a simplexity class, so `component` declares what this +# section configures. Without that declaration run management cannot tell a vendored process from +# any other instance config, because a vendored process shares no namespace with simplexity. +name: vendored_mess3 +component: generative_process +instance: + _target_: tests.vendored_process.adapter.build_vendored_mess3 + x: 0.15 + a: 0.6 + +base_vocab_size: ??? +bos_token: ??? +eos_token: null +vocab_size: ??? diff --git a/tests/end_to_end/configs/test_vendored_process.yaml b/tests/end_to_end/configs/test_vendored_process.yaml new file mode 100644 index 00000000..7f0c7bac --- /dev/null +++ b/tests/end_to_end/configs/test_vendored_process.yaml @@ -0,0 +1,17 @@ +# @package _global_ + +defaults: + - _self_ + - mlflow: databricks + - logging: mlflow_logger + - generative_process: vendored_mess3 + - predictive_model: tiny_transformer + +experiment_name: vendored_process_test +run_name: vendored_process_test_${now:%Y%m%d_%H%M%S} +device: auto +seed: 0 +logging_config_path: tests/end_to_end/configs/logging.ini +tags: + research_step: test + retention: temp diff --git a/tests/end_to_end/test_vendored_process_training.py b/tests/end_to_end/test_vendored_process_training.py new file mode 100644 index 00000000..19a0ba28 --- /dev/null +++ b/tests/end_to_end/test_vendored_process_training.py @@ -0,0 +1,128 @@ +"""Train a generative process vendored from generators, under `managed_run`. + +This is the acceptance test for instantiating a process that lives in the consumer's project +rather than in simplexity: the process code under `tests/vendored_process/generators_copy/` is a +copy, not an import, and neither it nor its adapter references +`simplexity.generative_processes`. +""" + +import math +from pathlib import Path +from typing import Any + +import jax +import jax.numpy as jnp +import pytest +import torch +from hydra import compose, initialize_config_dir +from omegaconf import DictConfig +from transformer_lens import HookedTransformer + +import simplexity +from simplexity.generative_processes.generative_process import ( + GenerativeProcess as GenerativeProcessBaseClass, +) +from simplexity.generative_processes.torch_generator import generate_data_batch +from simplexity.run_management.protocols import GenerativeProcessProtocol, LogSpaceGenerativeProcessProtocol +from tests.vendored_process.adapter import VendoredGhmmProcess + +CONFIG_DIR = str(Path(__file__).parent / "configs") +BATCH_SIZE = 4 +SEQUENCE_LEN = 8 +NUM_STEPS = 3 + + +def _compose(tmp_path: Path) -> DictConfig: + mlflow_uri = f"sqlite:///{(tmp_path / 'mlflow.db').absolute()}" + overrides = [f"mlflow.tracking_uri={mlflow_uri}", f"mlflow.registry_uri={mlflow_uri}"] + with initialize_config_dir(CONFIG_DIR, version_base="1.2"): + return compose(config_name="test_vendored_process.yaml", overrides=overrides) + + +@simplexity.managed_run(strict=False, verbose=False) +def _train_on_vendored_process(cfg: DictConfig, components: simplexity.Components) -> dict[str, Any]: + """Train the configured model on data generated by the vendored process.""" + generative_process = components.get_generative_process() + assert generative_process is not None, "vendored process was not instantiated" + + predictive_model = components.get_predictive_model() + assert isinstance(predictive_model, HookedTransformer) + + gen_states = jnp.repeat(generative_process.initial_state[None, :], BATCH_SIZE, axis=0) + optimizer = torch.optim.Adam(predictive_model.parameters(), lr=0.01) + loss_fn = torch.nn.CrossEntropyLoss() + + losses: list[float] = [] + inputs = torch.empty(0) + for step in range(NUM_STEPS): + gen_states, inputs, labels = generate_data_batch( + gen_states, + generative_process, + BATCH_SIZE, + SEQUENCE_LEN, + jax.random.key(step), + bos_token=cfg.generative_process.bos_token, + ) + outputs = predictive_model(inputs) + loss = loss_fn(outputs.reshape(-1, outputs.shape[-1]), labels.reshape(-1).long().to(outputs.device)) + optimizer.zero_grad() + loss.backward() + optimizer.step() + losses.append(float(loss.detach().item())) + + return { + "process": generative_process, + "losses": losses, + "inputs": inputs, + "d_vocab": predictive_model.cfg.d_vocab, + "resolved_vocab_size": cfg.generative_process.vocab_size, + "resolved_base_vocab_size": cfg.generative_process.base_vocab_size, + "resolved_bos_token": cfg.generative_process.bos_token, + } + + +@pytest.fixture(scope="module") +def training_result(tmp_path_factory: pytest.TempPathFactory) -> dict[str, Any]: + """Run one managed training run against the vendored process.""" + tmp_path = tmp_path_factory.mktemp("vendored_process") + return _train_on_vendored_process(_compose(tmp_path)) # pylint: disable=no-value-for-parameter + + +def test_vendored_process_is_instantiated(training_result: dict[str, Any]): + """The vendored process is discovered and instantiated, not silently skipped.""" + assert isinstance(training_result["process"], VendoredGhmmProcess) + + +def test_vendored_process_is_not_a_simplexity_process(training_result: dict[str, Any]): + """The process conforms structurally without inheriting from simplexity.""" + process = training_result["process"] + assert not isinstance(process, GenerativeProcessBaseClass) + assert isinstance(process, GenerativeProcessProtocol) + + +def test_vendored_process_lacks_log_space_operations(training_result: dict[str, Any]): + """Log-space operations are an optional extension that generators does not provide.""" + assert not isinstance(training_result["process"], LogSpaceGenerativeProcessProtocol) + + +def test_vocabulary_fields_resolve_from_the_vendored_process(training_result: dict[str, Any]): + """Vocabulary resolution reads the vendored process, so `d_vocab: ???` gets filled in.""" + assert training_result["resolved_base_vocab_size"] == 3 + assert training_result["resolved_bos_token"] == 3 + assert training_result["resolved_vocab_size"] == 4 + assert training_result["d_vocab"] == 4 + + +def test_generated_tokens_are_in_vocabulary(training_result: dict[str, Any]): + """Generated data has the expected shape and stays within the resolved vocabulary.""" + inputs = training_result["inputs"] + assert inputs.shape == (BATCH_SIZE, SEQUENCE_LEN) + assert int(inputs.min()) >= 0 + assert int(inputs.max()) < training_result["resolved_vocab_size"] + + +def test_training_steps_produce_finite_losses(training_result: dict[str, Any]): + """Training runs end to end on vendored-process data.""" + losses = training_result["losses"] + assert len(losses) == NUM_STEPS + assert all(loss > 0 and math.isfinite(loss) for loss in losses) diff --git a/tests/run_management/test_generative_process_discovery.py b/tests/run_management/test_generative_process_discovery.py new file mode 100644 index 00000000..449e97c8 --- /dev/null +++ b/tests/run_management/test_generative_process_discovery.py @@ -0,0 +1,119 @@ +"""Tests for discovering and instantiating generative processes, including foreign ones. + +"Foreign" means a process whose implementation lives outside `simplexity.generative_processes` — +typically vendored from generators into the consumer's project — so it cannot be recognized from +its import path. +""" + +import pytest +from omegaconf import DictConfig, OmegaConf + +from simplexity.exceptions import ConfigValidationError +from simplexity.run_management.run_management import ( + _instantiate_generative_process, + _setup_generative_processes, +) +from simplexity.utils.config_utils import get_instance_keys + +SIMPLEXITY_TARGET = "simplexity.generative_processes.builder.build_hidden_markov_model" +FOREIGN_TARGET = "tests.vendored_process.adapter.build_vendored_mess3" +FOREIGN_INSTANCE = {"_target_": FOREIGN_TARGET, "x": 0.15, "a": 0.6} +SIMPLEXITY_INSTANCE = {"_target_": SIMPLEXITY_TARGET, "process_name": "mess3", "process_params": {"x": 0.15, "a": 0.6}} + + +def _process_cfg(instance: dict, *, declare: bool) -> DictConfig: + section: dict = {"name": "process", "instance": instance} + if declare: + section["component"] = "generative_process" + section |= {"base_vocab_size": "???", "bos_token": "???", "eos_token": None, "vocab_size": "???"} + return OmegaConf.create({"generative_process": section}) + + +def _setup(cfg: DictConfig): + return _setup_generative_processes(cfg, get_instance_keys(cfg)) + + +def test_declared_foreign_process_is_instantiated() -> None: + """A declared foreign target is discovered even though its namespace is not simplexity.""" + processes = _setup(_process_cfg(FOREIGN_INSTANCE, declare=True)) + assert processes is not None + assert list(processes) == ["generative_process.instance"] + assert type(processes["generative_process.instance"]).__name__ == "VendoredGhmmProcess" + + +def test_declared_foreign_process_resolves_vocabulary_fields() -> None: + """Resolution reads the instantiated process, so it works for foreign processes too.""" + cfg = _process_cfg(FOREIGN_INSTANCE, declare=True) + _setup(cfg) + assert cfg.generative_process.base_vocab_size == 3 + assert cfg.generative_process.bos_token == 3 + assert cfg.generative_process.vocab_size == 4 + + +def test_simplexity_process_needs_no_declaration() -> None: + """The namespace fast path keeps every existing config working unchanged.""" + processes = _setup(_process_cfg(SIMPLEXITY_INSTANCE, declare=False)) + assert processes is not None + assert list(processes) == ["generative_process.instance"] + + +def test_declaring_a_simplexity_process_is_harmless() -> None: + """Declaration is additive, so configs may adopt it before the namespace path retires.""" + processes = _setup(_process_cfg(SIMPLEXITY_INSTANCE, declare=True)) + assert processes is not None + assert list(processes) == ["generative_process.instance"] + + +def test_undeclared_foreign_process_is_not_discovered() -> None: + """Without a declaration, a foreign target is indistinguishable from any other instance.""" + assert _setup(_process_cfg(FOREIGN_INSTANCE, declare=False)) is None + + +def test_config_without_a_process_yields_none() -> None: + """A config that genuinely configures no process is still not an error.""" + cfg = OmegaConf.create({"optimizer": {"instance": {"_target_": "torch.optim.Adam", "lr": 0.01}}}) + assert _setup(cfg) is None + + +def test_declared_process_with_an_invalid_config_raises() -> None: + """A section that declares a process must not be dropped silently. + + Before this behaviour existed, an unusable declaration left `generative_processes` as None, + which is indistinguishable from a config that configures no process at all. + """ + cfg = _process_cfg({"_target_": ""}, declare=True) + with pytest.raises(ConfigValidationError, match="declares generative processes at"): + _setup(cfg) + + +def test_declared_process_error_names_the_offending_key() -> None: + """The error identifies which section is at fault.""" + cfg = _process_cfg({"_target_": ""}, declare=True) + with pytest.raises(ConfigValidationError, match="generative_process.instance"): + _setup(cfg) + + +def test_undeclared_invalid_foreign_config_is_skipped_not_raised() -> None: + """Unclaimed sections stay skippable; only declared intent turns a skip into an error.""" + assert _setup(_process_cfg({"_target_": "torch.optim.Adam"}, declare=False)) is None + + +def test_instantiating_a_non_process_raises_naming_missing_members() -> None: + """A declared target that instantiates something else fails with an actionable message.""" + cfg = _process_cfg({"_target_": "builtins.dict"}, declare=True) + with pytest.raises(ConfigValidationError, match="which is not a generative process: missing"): + _instantiate_generative_process(cfg, "generative_process.instance") + + +def test_instantiating_a_non_process_lists_the_absent_operations() -> None: + """The error names the specific operations that are missing.""" + cfg = _process_cfg({"_target_": "builtins.dict"}, declare=True) + with pytest.raises(ConfigValidationError, match="emit_observation"): + _instantiate_generative_process(cfg, "generative_process.instance") + + +def test_instantiating_a_missing_instance_key_raises_key_error() -> None: + """An instance key that does not exist is a KeyError, not a silent None.""" + cfg = _process_cfg(FOREIGN_INSTANCE, declare=True) + with pytest.raises(KeyError): + _instantiate_generative_process(cfg, "generative_process.absent") diff --git a/tests/run_management/test_protocols.py b/tests/run_management/test_protocols.py new file mode 100644 index 00000000..158772a3 --- /dev/null +++ b/tests/run_management/test_protocols.py @@ -0,0 +1,136 @@ +"""Tests for the structural component contracts.""" + +import jax.numpy as jnp +import pytest + +from simplexity.generative_processes.builder import build_hidden_markov_model +from simplexity.generative_processes.generative_process import GenerativeProcess as GenerativeProcessBaseClass +from simplexity.run_management.protocols import ( + GenerativeProcessProtocol, + LogSpaceGenerativeProcessProtocol, + missing_generative_process_members, +) + +CORE_MEMBERS = { + "emit_observation", + "generate", + "initial_state", + "observation_probability_distribution", + "probability", + "transition_states", + "vocab_size", +} +LOG_SPACE_MEMBERS = {"log_observation_probability_distribution", "log_probability"} + + +class MinimalProcess: + """The smallest object satisfying the core protocol.""" + + vocab_size = 2 + initial_state = None + + def emit_observation(self, state, key): + """Emit an observation.""" + + def transition_states(self, state, obs): + """Transition states.""" + + def observation_probability_distribution(self, state): + """Observation distribution.""" + + def probability(self, observations): + """Sequence probability.""" + + def generate(self, state, key, sequence_len, return_all_states): + """Generate sequences.""" + + +class LogSpaceProcess(MinimalProcess): + """A core process that also provides log-space operations.""" + + def log_observation_probability_distribution(self, log_belief_state): + """Log observation distribution.""" + + def log_probability(self, observations): + """Log sequence probability.""" + + +def test_core_protocol_members_are_pinned() -> None: + """Guard the contract against accidental widening or narrowing. + + Asserted through the public helper rather than the protocol's internals, so the pin does not + depend on a CPython implementation detail. + """ + assert set(missing_generative_process_members(object())) == CORE_MEMBERS + + +def test_log_space_operations_are_not_in_the_core_contract() -> None: + """Log-space operations stay optional because only belief-state analysis needs them.""" + assert not LOG_SPACE_MEMBERS & CORE_MEMBERS + assert isinstance(MinimalProcess(), GenerativeProcessProtocol) + + +@pytest.mark.parametrize("member", sorted(LOG_SPACE_MEMBERS)) +def test_omitting_any_log_space_member_breaks_the_extension_only(member: str) -> None: + """Each log-space member is required by the extension and irrelevant to the core.""" + attrs = {name: getattr(LogSpaceProcess, name) for name in CORE_MEMBERS | LOG_SPACE_MEMBERS if name != member} + partial_process = type("PartialLogSpaceProcess", (), attrs)() + assert isinstance(partial_process, GenerativeProcessProtocol) + assert not isinstance(partial_process, LogSpaceGenerativeProcessProtocol) + + +def test_log_space_protocol_requires_the_core_members_too() -> None: + """The extension is the core contract plus log-space, not log-space alone.""" + attrs = {name: getattr(LogSpaceProcess, name) for name in LOG_SPACE_MEMBERS} + log_space_only = type("LogSpaceOnly", (), attrs)() + assert not isinstance(log_space_only, LogSpaceGenerativeProcessProtocol) + + +def test_minimal_object_satisfies_the_core_protocol() -> None: + """Conformance requires the operations, not a base class.""" + assert isinstance(MinimalProcess(), GenerativeProcessProtocol) + assert not isinstance(MinimalProcess(), GenerativeProcessBaseClass) + + +def test_core_process_does_not_satisfy_the_log_space_protocol() -> None: + """A process without log-space operations conforms to the core contract only.""" + assert not isinstance(MinimalProcess(), LogSpaceGenerativeProcessProtocol) + + +def test_log_space_process_satisfies_both_protocols() -> None: + """A process providing everything conforms to both.""" + process = LogSpaceProcess() + assert isinstance(process, GenerativeProcessProtocol) + assert isinstance(process, LogSpaceGenerativeProcessProtocol) + + +def test_simplexity_processes_satisfy_both_protocols() -> None: + """Relaxing the contract must not exclude the processes implemented in simplexity.""" + process = build_hidden_markov_model("mess3", {"x": 0.15, "a": 0.6}) + assert isinstance(process, GenerativeProcessProtocol) + assert isinstance(process, LogSpaceGenerativeProcessProtocol) + + +@pytest.mark.parametrize("member", sorted(CORE_MEMBERS)) +def test_omitting_any_core_member_breaks_conformance(member: str) -> None: + """Every member of the core contract is load-bearing, and the absent one is named.""" + attrs = {name: getattr(MinimalProcess, name) for name in CORE_MEMBERS if name != member} + partial_process = type("PartialProcess", (), attrs)() + assert not isinstance(partial_process, GenerativeProcessProtocol) + assert missing_generative_process_members(partial_process) == [member] + + +def test_missing_members_are_empty_for_a_conforming_process() -> None: + """A conforming process reports nothing missing.""" + assert missing_generative_process_members(MinimalProcess()) == [] + + +def test_missing_members_are_reported_sorted() -> None: + """Reporting which members are absent is what makes a rejected config actionable.""" + assert missing_generative_process_members(object()) == sorted(CORE_MEMBERS) + + +def test_missing_members_for_a_partially_conforming_object() -> None: + """Only the absent members are reported, not the ones provided.""" + half_process = type("HalfProcess", (), {"vocab_size": 2, "initial_state": jnp.zeros(2)})() + assert missing_generative_process_members(half_process) == sorted(CORE_MEMBERS - {"vocab_size", "initial_state"}) diff --git a/tests/structured_configs/test_generative_process_config.py b/tests/structured_configs/test_generative_process_config.py index fdcda804..3229e858 100644 --- a/tests/structured_configs/test_generative_process_config.py +++ b/tests/structured_configs/test_generative_process_config.py @@ -31,6 +31,9 @@ HiddenMarkovModelInstanceConfig, InstanceConfig, NonergodicHiddenMarkovModelBuilderInstanceConfig, + claims_generative_process_instance_key, + declares_generative_process, + declares_generative_process_instance_key, is_generalized_hidden_markov_model_builder_config, is_generalized_hidden_markov_model_builder_target, is_generalized_hidden_markov_model_config, @@ -1029,3 +1032,70 @@ def test_validate_generative_process_config_invalid_factored_process_configs( cfg = self._make_factored_process_cfg(structure_type, extra_instance_fields) # All these cases pass config validation (builder will validate later) validate_generative_process_config(cfg) + + +class TestForeignGenerativeProcessDeclaration: + """Tests for declaring a generative process implemented outside simplexity. + + Processes vendored from generators live in the consumer's namespace, so their configs cannot + be recognized from `_target_` and must declare themselves instead. + """ + + FOREIGN_TARGET = "my_project.processes.factory.build_process" + + def _make_cfg(self, target: str, component: str | None) -> DictConfig: + section: dict[str, Any] = {"name": "process", "instance": {"_target_": target}} + if component is not None: + section["component"] = component + section |= {"base_vocab_size": MISSING, "bos_token": MISSING, "eos_token": None, "vocab_size": MISSING} + return OmegaConf.create({"generative_process": section}) + + def test_declares_generative_process_true(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, "generative_process") + assert declares_generative_process(cfg.generative_process) + + def test_declares_generative_process_false_when_absent(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, None) + assert not declares_generative_process(cfg.generative_process) + + def test_declares_generative_process_false_for_another_component(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, "predictive_model") + assert not declares_generative_process(cfg.generative_process) + + def test_declares_instance_key(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, "generative_process") + assert declares_generative_process_instance_key(cfg, "generative_process.instance") + + def test_declares_instance_key_false_for_unknown_key(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, "generative_process") + assert not declares_generative_process_instance_key(cfg, "absent.instance") + + def test_claims_declared_foreign_instance_key(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, "generative_process") + assert claims_generative_process_instance_key(cfg, "generative_process.instance") + + def test_does_not_claim_undeclared_foreign_instance_key(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, None) + assert not claims_generative_process_instance_key(cfg, "generative_process.instance") + + def test_claims_simplexity_instance_key_without_declaration(self) -> None: + cfg = self._make_cfg("simplexity.generative_processes.hidden_markov_model.HiddenMarkovModel", None) + assert claims_generative_process_instance_key(cfg, "generative_process.instance") + + def test_does_not_claim_instance_key_without_a_target(self) -> None: + cfg = OmegaConf.create({"generative_process": {"instance": {"process_name": "mess3"}}}) + assert not claims_generative_process_instance_key(cfg, "generative_process.instance") + + def test_validation_accepts_a_declared_foreign_config(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, "generative_process") + validate_generative_process_config(cfg.generative_process) + + def test_validation_rejects_an_undeclared_foreign_config(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, None) + with pytest.raises(ConfigValidationError, match="must be a generative process target"): + validate_generative_process_config(cfg.generative_process) + + def test_rejection_message_explains_the_declaration(self) -> None: + cfg = self._make_cfg(self.FOREIGN_TARGET, None) + with pytest.raises(ConfigValidationError, match="component: generative_process"): + validate_generative_process_config(cfg.generative_process) diff --git a/tests/utils/test_config_utils.py b/tests/utils/test_config_utils.py index 798ebecf..52069a9d 100644 --- a/tests/utils/test_config_utils.py +++ b/tests/utils/test_config_utils.py @@ -9,8 +9,10 @@ TARGET, dynamic_resolve, filter_instance_keys, + filter_instance_keys_by, get_config, get_instance_keys, + get_instance_target, typed_instantiate, ) @@ -372,3 +374,61 @@ def test_typed_instantiate_with_string_expected_type() -> None: obj = typed_instantiate(cfg, "builtins.str") assert obj == "42" assert isinstance(obj, str) + + +def test_get_instance_target() -> None: + """Test reading an instance key's target.""" + cfg = DictConfig({"instance": DictConfig({TARGET: "some_callable"})}) + assert get_instance_target(cfg, "instance") == "some_callable" + + +def test_get_instance_target_missing_key() -> None: + """Test reading the target of an absent instance key.""" + cfg = DictConfig({"instance": DictConfig({TARGET: "some_callable"})}) + assert get_instance_target(cfg, "absent") is None + + +def test_get_instance_target_non_string() -> None: + """Test reading a target that is not a string.""" + cfg = DictConfig({"instance": DictConfig({TARGET: None})}) + assert get_instance_target(cfg, "instance") is None + + +def test_filter_instance_keys_by_consults_the_whole_config() -> None: + """Test filtering by a predicate that reads beyond the instance's target. + + Components whose implementations may live outside simplexity cannot be identified from the + target string alone, so the predicate receives the config and the instance key. + """ + cfg = DictConfig( + { + "component1": DictConfig({"instance": DictConfig({TARGET: "foreign_callable"}), "claimed": True}), + "component2": DictConfig({"instance": DictConfig({TARGET: "foreign_callable"}), "claimed": False}), + } + ) + + def claims_fn(cfg: DictConfig, instance_key: str) -> bool: + section = OmegaConf.select(cfg, instance_key.rsplit(".", 1)[0]) + return bool(section and section.get("claimed", False)) + + instance_keys = ["component1.instance", "component2.instance"] + assert filter_instance_keys_by(cfg, instance_keys, claims_fn) == ["component1.instance"] + + +def test_filter_instance_keys_by_applies_validation() -> None: + """Test that a claimed instance key still has to validate.""" + cfg = DictConfig({"component1": DictConfig({"instance": DictConfig({TARGET: "some_callable"})})}) + + def claims_fn(_cfg: DictConfig, _instance_key: str) -> bool: + return True + + def validate_fn(cfg: DictConfig) -> None: + raise ConfigValidationError("invalid") + + assert filter_instance_keys_by(cfg, ["component1.instance"], claims_fn, validate_fn, "test") == [] + + +def test_filter_instance_keys_by_rejects_everything_when_predicate_is_false() -> None: + """Test filtering when the predicate claims nothing.""" + cfg = DictConfig({"component1": DictConfig({"instance": DictConfig({TARGET: "some_callable"})})}) + assert filter_instance_keys_by(cfg, ["component1.instance"], lambda _cfg, _key: False) == [] diff --git a/tests/vendored_process/__init__.py b/tests/vendored_process/__init__.py new file mode 100644 index 00000000..84bcbbca --- /dev/null +++ b/tests/vendored_process/__init__.py @@ -0,0 +1 @@ +"""Vendored generative process fixtures and their simplexity adapter.""" diff --git a/tests/vendored_process/adapter.py b/tests/vendored_process/adapter.py new file mode 100644 index 00000000..4db8298d --- /dev/null +++ b/tests/vendored_process/adapter.py @@ -0,0 +1,90 @@ +"""Project-local adapter making vendored generators modules runnable by simplexity. + +This is the worked example the migration guide points at. It is the only file a consumer writes: +the vendored modules stay verbatim, and simplexity is not modified at all. + +Two things need bridging. Generators exposes module-level functions over a `NamedTuple` of process +data, so the adapter binds that data once and exposes the operations as methods. And generators' +`generate` returns only the final state, whereas simplexity's training path also wants every +intermediate belief state, so the adapter re-scans rather than calling it. +""" + +from typing import Any + +import chex +import equinox as eqx +import jax +import jax.numpy as jnp + +from tests.vendored_process.generators_copy import ghmm_process +from tests.vendored_process.generators_copy.classical import mess + + +class VendoredGhmmProcess(eqx.Module): + """A generalized hidden Markov model backed by vendored generators code. + + Structurally satisfies `simplexity.run_management.protocols.GenerativeProcessProtocol` without + importing or subclassing anything from `simplexity.generative_processes`. + """ + + data: ghmm_process.Data + + @property + def vocab_size(self) -> int: + """The number of observations that can be emitted by the generative process.""" + return int(self.data.Ts.shape[0]) + + @property + def initial_state(self) -> jax.Array: + """The stationary belief state that generators' `init` derived from the process.""" + return self.data.eta_0 + + def emit_observation(self, state: jax.Array, key: chex.PRNGKey) -> chex.Array: + """Emit an observation based on the state of the generative process.""" + return ghmm_process.sample(self.data, state, key) + + def transition_states(self, state: jax.Array, obs: chex.Array) -> jax.Array: + """Evolve the state of the generative process based on the observation.""" + return ghmm_process.update(self.data, state, jnp.asarray(obs)) + + def observation_probability_distribution(self, state: jax.Array) -> jax.Array: + """Compute the distribution over observations that can be emitted from a state.""" + return ghmm_process.obs_dist(self.data, state) + + def probability(self, observations: jax.Array) -> jax.Array: + """Compute the probability of the process generating a sequence of observations.""" + return ghmm_process.seq_prob(self.data, observations) + + @eqx.filter_vmap(in_axes=(None, 0, 0, None, None)) + def generate( + self, state: jax.Array, key: chex.PRNGKey, sequence_len: int, return_all_states: bool + ) -> tuple[Any, chex.Array]: + """Generate a batch of observation sequences, optionally returning every belief state.""" + keys = jax.random.split(key, sequence_len) + + def step(state: jax.Array, key: chex.PRNGKey) -> tuple[jax.Array, tuple[jax.Array, chex.Array]]: + obs = self.emit_observation(state, key) + return self.transition_states(state, obs), (state, obs) + + final_state, (states, observations) = jax.lax.scan(step, state, keys) + if return_all_states: + return states, observations + return final_state, observations + + +def build_vendored_mess3(x: float, a: float, num_states: int = 3) -> VendoredGhmmProcess: + """Build a vendored mess3 process. + + Hydra targets this rather than the class, because assembling a process from generators modules + is code: transition matrices are constructed, passed through `init`, and for composite + processes bound to consumer-chosen encode/decode callables. None of that fits in YAML. + + Args: + x: Transition probability to each other state. + a: Emission probability corresponding to the previous state. + num_states: Number of hidden states. + + Returns: + A process satisfying simplexity's generative process protocol. + """ + return VendoredGhmmProcess(data=ghmm_process.init(mess(x, a, num_states))) diff --git a/tests/vendored_process/generators_copy/README.md b/tests/vendored_process/generators_copy/README.md new file mode 100644 index 00000000..ba327232 --- /dev/null +++ b/tests/vendored_process/generators_copy/README.md @@ -0,0 +1,19 @@ +# Vendored generators modules + +Verbatim copies of [generators](https://github.com/ealt/generators) at commit `b1242ea`, +apart from rewriting the module-level import to this package's path: + +| File | Upstream path | +| --- | --- | +| `ghmm_process.py` | `generators/ghmm/process.py` | +| `utils.py` | `generators/utils.py` | + +They are checked in, rather than imported, because that is generators' distribution model — the +copy is the point. It also keeps this test suite independent of generators being installed. + +`../adapter.py` binds these module-level functions to +`simplexity.run_management.protocols.GenerativeProcessProtocol`, and +`../../end_to_end/test_vendored_process_training.py` trains against it under `@managed_run()`. + +To pick up upstream improvements, diff against the current upstream module and re-copy. Do not add +simplexity-specific behaviour here; that belongs in the adapter. diff --git a/tests/vendored_process/generators_copy/__init__.py b/tests/vendored_process/generators_copy/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/vendored_process/generators_copy/classical.py b/tests/vendored_process/generators_copy/classical.py new file mode 100644 index 00000000..529b19a8 --- /dev/null +++ b/tests/vendored_process/generators_copy/classical.py @@ -0,0 +1,75 @@ +import jax +import jax.numpy as jnp + + +def _mess_trans(x: float, s: int) -> jax.Array: + r"""State transition matrix for Mess process. + + Args: + x: Transition probability to each other state. + x = P(s_t = s' \forall s' \in {0, ..., s-1} \ s') + y: Probability of staying in the same state. + y = P(s_t = s_t-1) + s: Number of states + + Returns: + State transition matrix + T[i, j] = P(s_t = i | s_{t-1} = j) + T[i, i] = y + T[i, j] = x for i != j + """ + assert x >= 0 + assert x <= 1 + y = 1 - (s - 1) * x + assert y >= 0 + assert y <= 1 + return (x * (jnp.ones((s, s)) - jnp.eye(s))) + (y * jnp.eye(s)) + + +def _mess_emit(a: float, s: int) -> jax.Array: + r"""Emission matrix for Mess process. + + Args: + a: Emission probability corresponding to the previous state. + a = P(x_t = s_t-1) + b: Emission probabilities for each other value. + b = P(x_t = s' \forall s' \in S_{t-1}) + s: Number of states + rd: Ratio difference + If the vocab size to state size ratio, V : S + is reduced, num : denom (V / S = num / denom) + rd := num - denom + + Returns: + Emission matrix + Tv[i, j] = P(x_t = j | s_t = i) + Tv[i, i] = a + Tv[i, j] = b for i != j + """ + assert a >= 0 + assert a <= 1 + b = (1 - a) / (s - 1) + + return (a * jnp.eye(s)) + (b * (jnp.ones((s, s)) - jnp.eye(s))) + + +def mess(x: float, a: float, s: int) -> jax.Array: + r"""Mess process. + + Args: + x: Transition probability to each other state. + a: Emission probability corresponding to the previous state. + s: Number of states + rd: Ratio difference + + Returns: + Transition matrix + Ts[o, j, k] = P(x_t=o, s_t=k | s_{t-1}=j) + """ + assert s > 0 + T_trans = _mess_trans(x, s) + T_emit = _mess_emit(a, s) + + # Ts[o, j, k] = P(x_t=o, s_t=k | s_{t-1}=j) = Tv[k, o] * T[k, j] + inner = T_emit[:, :, None] * T_trans[None, :, :] + return jnp.transpose(inner, (0, 2, 1)) diff --git a/tests/vendored_process/generators_copy/ghmm_process.py b/tests/vendored_process/generators_copy/ghmm_process.py new file mode 100644 index 00000000..3fb4c6f9 --- /dev/null +++ b/tests/vendored_process/generators_copy/ghmm_process.py @@ -0,0 +1,66 @@ +from typing import NamedTuple + +import jax +import jax.numpy as jnp + +from tests.vendored_process.generators_copy.utils import principal_ev + + +class Data(NamedTuple): + Ts: jax.Array + eta_0: jax.Array + w: jax.Array + + +def validate(Ts: jax.Array) -> bool: + if len(Ts.shape) != 3: + return False + if any(dim == 0 for dim in Ts.shape): + return False + if Ts.shape[1] != Ts.shape[2]: + return False + if not jnp.all(jnp.isfinite(Ts)): + return False + T = jnp.sum(Ts, axis=0) + norm = jnp.linalg.norm(T, ord=jnp.inf) + return bool(jnp.isclose(norm, 1)) + + +def init(Ts: jax.Array) -> Data: + T = Ts.sum(axis=0) + w = principal_ev(T) + eta_0 = principal_ev(T.T) + eta_0 /= eta_0 @ w + return Data(Ts=Ts, eta_0=eta_0, w=w) + + +def obs_dist(data: Data, eta: jax.Array) -> jax.Array: + return eta @ data.Ts @ data.w + + +def sample(data: Data, eta: jax.Array, key: jax.Array) -> jax.Array: + probs = obs_dist(data, eta) + logits = jnp.where(probs > 0, jnp.log(probs), -jnp.inf) + return jax.random.categorical(key, logits) + + +def update(data: Data, eta: jax.Array, x: jax.Array) -> jax.Array: + eta = eta @ data.Ts[x] + return eta / (eta @ data.w) + + +def generate(data: Data, eta: jax.Array, keys: jax.Array) -> tuple[jax.Array, jax.Array]: + def fn(eta: jax.Array, key: jax.Array) -> tuple[jax.Array, jax.Array]: + x = sample(data, eta, key) + eta = update(data, eta, x) + return eta, x + + return jax.lax.scan(fn, eta, keys) + + +def seq_prob(data: Data, xs: jax.Array) -> jax.Array: + def fn(eta, x): + return eta @ data.Ts[x], None + + eta, _ = jax.lax.scan(fn, init=data.eta_0, xs=xs) + return eta @ data.w diff --git a/tests/vendored_process/generators_copy/utils.py b/tests/vendored_process/generators_copy/utils.py new file mode 100644 index 00000000..1f578656 --- /dev/null +++ b/tests/vendored_process/generators_copy/utils.py @@ -0,0 +1,34 @@ +import jax +import jax.numpy as jnp + + +def principal_ev(T: jax.Array) -> jax.Array: + eigenvalues, eigenvectors = jnp.linalg.eig(T) + i = jnp.argmax(jnp.abs(eigenvalues)) + vector = jnp.real(eigenvectors[:, i]) + sign = jnp.where(jnp.sum(vector) < 0, -1.0, 1.0) + vector = vector * sign + return vector / jnp.mean(vector) + + +def pad(arr: jax.Array, target_shape: tuple[int, ...]) -> jax.Array: + pad_width = [(0, target - current) for target, current in zip(target_shape, arr.shape, strict=True)] + return jnp.pad(arr, pad_width, mode="constant", constant_values=0) + + +def stack(arrs: list[jax.Array]) -> jax.Array: + shapes = [arr.shape for arr in arrs] + max_shape = tuple(max(shape_i) for shape_i in zip(*shapes, strict=True)) + return jnp.stack([pad(arr, max_shape) for arr in arrs]) + + +def mixed_radix_weights(Vs: jax.Array) -> jax.Array: + return jnp.roll(jnp.cumprod(Vs), 1).at[0].set(1) + + +def mixed_radix_encode(x_factors: jax.Array, *, weights: jax.Array) -> jax.Array: + return jnp.sum(x_factors * weights) + + +def mixed_radix_decode(x: jax.Array, *, Vs: jax.Array, weights: jax.Array) -> jax.Array: + return (x // weights) % Vs