Instantiate generative processes by capability, not by namespace

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()`.

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

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,
+)

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:

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,

simplexity/run_management/components.py

@@ -13,18 +13,18 @@
 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
 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")