diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 94227768..d12aaddb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -40,3 +40,17 @@ jobs: python-version: ${{ matrix.python-version }} - run: pip install -e ".[dev,data,diagnostics,formulas,repl,lsp]" - run: python -m pytest tests/ -x -q + + # The default addopts deselect `slow`, so the gallery sweep and the + # recovery benchmarks never run in the `test` job. They gate here + # instead: an example whose prose or code stops matching the library + # fails CI rather than rotting unnoticed. + slow: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.14" + - run: pip install -e ".[dev,data,diagnostics,formulas,repl,lsp]" + - run: python -m pytest tests/ -m slow -q diff --git a/CHANGELOG.md b/CHANGELOG.md index 186b161a..8a6fc19d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,34 @@ All notable changes to the quivers library are documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/), and this project adheres to [Semantic Versioning](https://semver.org/). +## [0.16.0] - 2026-07-15 + +### Changed + +- **A kernel's parameter source defaults to a linear map.** A morphism declared `f : X -> Y ~ Family` over continuous spaces reads its distribution parameters from a single `nn.Linear`, which is how its arrow reads. A model whose parameters depend nonlinearly on the input requests that with `[param_source=mlp]`, so the fact appears in the source instead of in a default: previously every such kernel carried a two-hidden-layer tanh MLP, and a `Real 1 -> Real 1` kernel silently held 4418 weights. The gallery is now what it says it is. `linear_gaussian_ssm` is linear, and an LSTM, GRU or vanilla-RNN gate is one matrix under a cell whose nonlinearity is the sigmoid its program body applies. `deep_markov`, the VAE's encoder and decoder bodies, the transformer and seq2seq feed-forward sublayers, and the recurrent cells that carry their own tanh now name the MLP they rely on. +- **Morphism options are checked against the role that reads them.** Options were checked against the union of every role's keys and then read per role, so an option belonging to another role passed validation and was dropped in silence. `[scale=]` on a family-backed kernel is the case that bit: nothing reads it there, and all 34 uses across the examples configured an init that never existed. Each role now declares the keys its lowering consumes, and a key outside that set is rejected with an error naming the role that does read it. +- **A default coefficient prior is autoscaled to its column.** A column enters a formula's linear predictor as `beta * column`, so a prior on the coefficient alone states the scale of its contribution, and one fixed `Normal(0, 5)` therefore means something different for every column. The default's scale is now divided by the column's root-mean-square, which states it in contribution space; an explicit per-name prior is emitted as written. Coefficients stay on their own column's scale, so nothing is transformed back. + +### Added + +- **Inline distributions take any number of vector-typed parameters.** A family may declare vector parameters in any position, and the stacked input is split at each one's declared dim rather than by a single trailing parameter eating the remaining columns. `MixtureNormal(weights, locations, scales)`, whose three per-component vectors are each shared across the response plate, is expressible. +- **`[param_source=(...)]` carries its arguments.** The grammar has always parsed the call form and `param_source_from_option` has always read parenthesised hidden widths, but the compiler decoded the option as a bare identifier and raised on anything else, so `[param_source=mlp(16, 8)]` died at compile time and the width parser was unreachable from QVR. +- **`hidden_dim` takes a sequence of widths.** One entry per hidden layer, so `[param_source=mlp, hidden_dim=[64, 32, 16]]` builds those three; a bare `hidden_dim=64` is the one-layer case. A single number could say how wide an MLP's layers are but not how many of them there were. + +### Fixed + +- **Every family that has a parameter source can select it.** `[param_source=...]` is accepted by the compiler for any family-backed kernel, and exactly one of the 28 conditional families could receive it: `_make_continuous_morphism` passed the keyword into a constructor that only `ConditionalNormal` declared, so the other 27 died on a `TypeError` raised inside `__init__`, with no line or column on it. Every family whose parameters come from a source now takes `param_source` and `param_source_option` and hands them to the one seam that builds it. The four whose parameters do not (`Horseshoe`, `GaussianProcess`, `Independent`, `Transformed`) refuse the option at compile time, naming the family. +- **`hidden_dim` is read or refused, not dropped.** It was ignored whenever `param_source` was given, so `[param_source=mlp, hidden_dim=32]` silently built the default width. It now reaches the source that has layers to apply it to, and a width aimed at a source with none is an error: asking a single matrix how wide its hidden layers should be has no answer, and silence would read as one. +- **A scalar response is scored against its codomain's event axis.** A conditional family's parameters for a `d`-dimensional codomain arrive as `(N, d)`. An `(N,)` response against a `d = 1` codomain, which is the shape the response placeholder implies, broadcast to `(N, N)`, and because the observe step sums its score the result was a finite log-joint wrong by a factor of `N`. The axis is restored, and a response that cannot carry the codomain's event shape is rejected rather than broadcast into one. An inline family keeps the plate layout its parameters already align with. +- **An unsupplied `via` fibration index names itself.** A grouped `marginalize` reads its index from the observations dict like any covariate, so omitting it is a user error; it surfaced as a bare `KeyError`. +- **`quivers.__version__` reads the installed distribution's metadata.** It was a literal, and it had drifted a release behind. + +### Documentation + +- **The Bayesian neural network is a neural network.** The example composed four latent morphisms under an algebra whose composition is matrix multiplication, so with no pointwise nonlinearity the chain collapsed to a single linear map; its fit drove a learnable input morphism at pure noise and reported a residual worse than predicting zero, its NUTS block named an `x` and an `observations` no earlier block defined, over a composite carrying no likelihood to sample, and its walkthrough placed matrix-normal priors through a syntax that does not parse. It is now a Bayesian MLP for nonlinear regression whose conditional-Normal kernel emits both a mean and a log-scale, fit to a sine wave that no linear model can represent and scored against the best least-squares line under a matching Gaussian. +- **The Gaussian mixture is scored per row.** The example grouped a per-row component assignment under a `marginalize` block keyed by a fibration index, which read as an item-grouped mixture and carried a discrete latent no guide could biject. `MixtureNormal` integrates the assignment out in closed form. +- **`ParamSource` is documented.** Its eight sources and the `[param_source=...]` option had no API page, despite being where a kernel's dependence on its input is computed. + ## [0.15.0] - 2026-07-02 ### Changed diff --git a/docs/api/continuous/param_source.md b/docs/api/continuous/param_source.md new file mode 100644 index 00000000..f9e09077 --- /dev/null +++ b/docs/api/continuous/param_source.md @@ -0,0 +1,42 @@ +# Parameter Sources + +`quivers.continuous.param_source` provides the map from a conditional +family's input to its distribution parameters. A morphism declared +`~ Family` over a continuous domain is a Kleisli arrow whose +parameters are produced by a `ParamSource`, so the source is where a +kernel's dependence on its input is computed, and where any +nonlinearity in that dependence lives. + +The concrete sources cover the standard architectures: `LinearSource` +is a single `nn.Linear`; `MLPSource` is a multi-layer perceptron with +configurable widths and activation; `AttentionSource` is a +self-attention head; `LookupSource` and `EmbeddingSource` handle +discrete domains; `IdentitySource`, `FunctionSource`, and +`ComposeSource` cover pass-through, a fixed callable, and composition +of two sources. + +`make_param_source` is the factory the families call, and +`param_source_from_option` parses the DSL's `[param_source=]` +morphism option. The default for a continuous domain is `LinearSource`, +so a kernel is linear unless it asks for something else; a `SetObject` +domain always uses `LookupSource` regardless of the requested kind. The +[Bayesian Neural Network](../../examples/bnn.md) example selects the +MLP source explicitly and relies on it for its nonlinearity. + +The hidden widths come either from the option's arguments or from +`hidden_dim`, one width per hidden layer: + + +```qvr +morphism f : X -> Y [param_source=mlp] ~ Normal # (64, 64) +morphism f : X -> Y [param_source=mlp(64, 32)] ~ Normal # (64, 32) +morphism f : X -> Y [param_source=mlp, hidden_dim=[64, 32]] ~ Normal # (64, 32) +morphism f : X -> Y [param_source=mlp, hidden_dim=64] ~ Normal # (64,) +``` + +A width given to a source with no hidden layers to apply it to is an +error rather than a silent no-op, and so is `param_source` on a family +whose parameters do not come from a source at all (`Horseshoe`, +`GaussianProcess`, `Independent`, `Transformed`). + +::: quivers.continuous.param_source diff --git a/docs/developer/changelog.md b/docs/developer/changelog.md index 186b161a..8a6fc19d 100644 --- a/docs/developer/changelog.md +++ b/docs/developer/changelog.md @@ -4,6 +4,34 @@ All notable changes to the quivers library are documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/), and this project adheres to [Semantic Versioning](https://semver.org/). +## [0.16.0] - 2026-07-15 + +### Changed + +- **A kernel's parameter source defaults to a linear map.** A morphism declared `f : X -> Y ~ Family` over continuous spaces reads its distribution parameters from a single `nn.Linear`, which is how its arrow reads. A model whose parameters depend nonlinearly on the input requests that with `[param_source=mlp]`, so the fact appears in the source instead of in a default: previously every such kernel carried a two-hidden-layer tanh MLP, and a `Real 1 -> Real 1` kernel silently held 4418 weights. The gallery is now what it says it is. `linear_gaussian_ssm` is linear, and an LSTM, GRU or vanilla-RNN gate is one matrix under a cell whose nonlinearity is the sigmoid its program body applies. `deep_markov`, the VAE's encoder and decoder bodies, the transformer and seq2seq feed-forward sublayers, and the recurrent cells that carry their own tanh now name the MLP they rely on. +- **Morphism options are checked against the role that reads them.** Options were checked against the union of every role's keys and then read per role, so an option belonging to another role passed validation and was dropped in silence. `[scale=]` on a family-backed kernel is the case that bit: nothing reads it there, and all 34 uses across the examples configured an init that never existed. Each role now declares the keys its lowering consumes, and a key outside that set is rejected with an error naming the role that does read it. +- **A default coefficient prior is autoscaled to its column.** A column enters a formula's linear predictor as `beta * column`, so a prior on the coefficient alone states the scale of its contribution, and one fixed `Normal(0, 5)` therefore means something different for every column. The default's scale is now divided by the column's root-mean-square, which states it in contribution space; an explicit per-name prior is emitted as written. Coefficients stay on their own column's scale, so nothing is transformed back. + +### Added + +- **Inline distributions take any number of vector-typed parameters.** A family may declare vector parameters in any position, and the stacked input is split at each one's declared dim rather than by a single trailing parameter eating the remaining columns. `MixtureNormal(weights, locations, scales)`, whose three per-component vectors are each shared across the response plate, is expressible. +- **`[param_source=(...)]` carries its arguments.** The grammar has always parsed the call form and `param_source_from_option` has always read parenthesised hidden widths, but the compiler decoded the option as a bare identifier and raised on anything else, so `[param_source=mlp(16, 8)]` died at compile time and the width parser was unreachable from QVR. +- **`hidden_dim` takes a sequence of widths.** One entry per hidden layer, so `[param_source=mlp, hidden_dim=[64, 32, 16]]` builds those three; a bare `hidden_dim=64` is the one-layer case. A single number could say how wide an MLP's layers are but not how many of them there were. + +### Fixed + +- **Every family that has a parameter source can select it.** `[param_source=...]` is accepted by the compiler for any family-backed kernel, and exactly one of the 28 conditional families could receive it: `_make_continuous_morphism` passed the keyword into a constructor that only `ConditionalNormal` declared, so the other 27 died on a `TypeError` raised inside `__init__`, with no line or column on it. Every family whose parameters come from a source now takes `param_source` and `param_source_option` and hands them to the one seam that builds it. The four whose parameters do not (`Horseshoe`, `GaussianProcess`, `Independent`, `Transformed`) refuse the option at compile time, naming the family. +- **`hidden_dim` is read or refused, not dropped.** It was ignored whenever `param_source` was given, so `[param_source=mlp, hidden_dim=32]` silently built the default width. It now reaches the source that has layers to apply it to, and a width aimed at a source with none is an error: asking a single matrix how wide its hidden layers should be has no answer, and silence would read as one. +- **A scalar response is scored against its codomain's event axis.** A conditional family's parameters for a `d`-dimensional codomain arrive as `(N, d)`. An `(N,)` response against a `d = 1` codomain, which is the shape the response placeholder implies, broadcast to `(N, N)`, and because the observe step sums its score the result was a finite log-joint wrong by a factor of `N`. The axis is restored, and a response that cannot carry the codomain's event shape is rejected rather than broadcast into one. An inline family keeps the plate layout its parameters already align with. +- **An unsupplied `via` fibration index names itself.** A grouped `marginalize` reads its index from the observations dict like any covariate, so omitting it is a user error; it surfaced as a bare `KeyError`. +- **`quivers.__version__` reads the installed distribution's metadata.** It was a literal, and it had drifted a release behind. + +### Documentation + +- **The Bayesian neural network is a neural network.** The example composed four latent morphisms under an algebra whose composition is matrix multiplication, so with no pointwise nonlinearity the chain collapsed to a single linear map; its fit drove a learnable input morphism at pure noise and reported a residual worse than predicting zero, its NUTS block named an `x` and an `observations` no earlier block defined, over a composite carrying no likelihood to sample, and its walkthrough placed matrix-normal priors through a syntax that does not parse. It is now a Bayesian MLP for nonlinear regression whose conditional-Normal kernel emits both a mean and a log-scale, fit to a sine wave that no linear model can represent and scored against the best least-squares line under a matching Gaussian. +- **The Gaussian mixture is scored per row.** The example grouped a per-row component assignment under a `marginalize` block keyed by a fibration index, which read as an item-grouped mixture and carried a discrete latent no guide could biject. `MixtureNormal` integrates the assignment out in closed form. +- **`ParamSource` is documented.** Its eight sources and the `[param_source=...]` option had no API page, despite being where a kernel's dependence on its input is computed. + ## [0.15.0] - 2026-07-02 ### Changed diff --git a/docs/examples/bidirectional-rnn-lm.md b/docs/examples/bidirectional-rnn-lm.md index 22fc674a..17fc34fe 100644 --- a/docs/examples/bidirectional-rnn-lm.md +++ b/docs/examples/bidirectional-rnn-lm.md @@ -12,9 +12,9 @@ object Embedded, FwdHidden, BwdHidden : Real 64 object Combined : Real 128 morphism tok_embed : Token -> Embedded [role=embed] -morphism fwd_cell : Embedded * FwdHidden -> FwdHidden [scale=0.1] ~ Normal -morphism bwd_cell : Embedded * BwdHidden -> BwdHidden [scale=0.1] ~ Normal -morphism combine : Combined -> Combined [scale=0.1] ~ Normal +morphism fwd_cell : Embedded * FwdHidden -> FwdHidden [param_source=mlp] ~ Normal +morphism bwd_cell : Embedded * BwdHidden -> BwdHidden [param_source=mlp] ~ Normal +morphism combine : Combined -> Combined [param_source=mlp] ~ Normal morphism lm_head : Combined -> Token ~ Categorical define forward_path = tok_embed >> scan(fwd_cell) diff --git a/docs/examples/bnn.md b/docs/examples/bnn.md index 301e744d..ff1ce107 100644 --- a/docs/examples/bnn.md +++ b/docs/examples/bnn.md @@ -2,131 +2,140 @@ ## Overview -A [Bayesian neural network](https://en.wikipedia.org/wiki/Bayesian_neural_network) ([MacKay 1992](https://doi.org/10.1162/neco.1992.4.3.448)) puts a prior over every weight and recovers a posterior over weights via SVI or MCMC, giving calibrated predictive uncertainty far from the training data. This example builds a three-layer Bayesian *linear* network whose per-layer weight matrices carry [matrix-normal](https://en.wikipedia.org/wiki/Matrix_normal_distribution) priors. The forward pass is a pure categorical composition of [`LatentMorphism`](../api/core/morphisms.md) factors: +A [Bayesian neural network](https://en.wikipedia.org/wiki/Bayesian_neural_network) ([MacKay 1992](https://doi.org/10.1162/neco.1992.4.3.448)) puts a prior over every weight and recovers a posterior over weights, giving calibrated predictive uncertainty far from the training data. This example fits a [multi-layer perceptron](https://en.wikipedia.org/wiki/Multilayer_perceptron) to a nonlinear regression target: the response is [Normal](https://en.wikipedia.org/wiki/Normal_distribution) about a learned function of the input, $$ -\mathsf{Item} \xrightarrow{X} \mathsf{H}_\text{in} \xrightarrow{W_1} \mathsf{H}_1 \xrightarrow{W_2} \mathsf{H}_2 \xrightarrow{W_3} \mathsf{H}_\text{out}. +\mu(x),\, \sigma(x) \;=\; \mathrm{MLP}(x), \qquad y_n \;\sim\; \mathcal{N}\!\bigl(\mu(x_n),\, \sigma(x_n)\bigr). $$ -Each $W_l$ is a learnable morphism whose prior is matrix-normal over its (in, out) [Kronecker covariance](https://en.wikipedia.org/wiki/Kronecker_product), the natural prior for a linear layer. Under `composition real [level=algebra]` the composition `X >> W_1 >> W_2 >> W_3` is a real-valued matmul stack. +The network emits both the mean and the log-scale, so the model is [heteroscedastic](https://en.wikipedia.org/wiki/Heteroscedasticity): the predictive spread varies with the input rather than being pinned to a single global noise level. + +The nonlinearity lives in the morphism's parameter network. A morphism declared `~ Normal` over continuous spaces is a [Kleisli arrow](https://en.wikipedia.org/wiki/Kleisli_category) whose distribution parameters are produced from its input by a [`ParamSource`](../api/continuous/param_source.md#quivers.continuous.param_source.ParamSource), and `[param_source=mlp, hidden_dim=64]` selects an [`MLPSource`](../api/continuous/param_source.md#quivers.continuous.param_source.MLPSource): two hidden layers of width 64 with tanh activations between them. That is where the model departs from a linear map, and it is the reason this example can fit a curve that no linear model can. ## QVR Source ```qvr -composition real [level=algebra] - -object Item : FinSet 200 -object H_in : FinSet 4 -object H1 : FinSet 32 -object H2 : FinSet 16 -object H_out : FinSet 2 - -morphism X : Item -> H_in [role=latent] +object Feature : Real 1 +object Target : Real 1 +object Resp : FinSet 200 -morphism W_1 : H_in -> H1 [role=latent] -morphism W_2 : H1 -> H2 [role=latent] -morphism W_3 : H2 -> H_out [role=latent] +morphism net : Feature -> Target [param_source=mlp, hidden_dim=64] ~ Normal -define bnn = X >> W_1 >> W_2 >> W_3 +program bnn : Resp -> Resp + observe y : Resp <- net(x) + return y export bnn ``` ## Walkthrough -The three weight declarations +`object Feature : Real 1` and `object Target : Real 1` declare the input and response as one-dimensional continuous spaces; `object Resp : FinSet 200` is the discrete plate indexing the rows. The plate is what the `observe` step reduces over, so it is the program's domain. - -```qvr -morphism W_1 : H_in -> H1 [role=latent] ~ MatrixNormal(0.0, 1.0, 1.0) over (dom, cod) -morphism W_2 : H1 -> H2 [role=latent] ~ MatrixNormal(0.0, 1.0, 1.0) over (dom, cod) -morphism W_3 : H2 -> H_out [role=latent] ~ MatrixNormal(0.0, 1.0, 1.0) over (dom, cod) -``` +`morphism net : Feature -> Target [param_source=mlp, hidden_dim=64] ~ Normal` declares the network. The `~ Normal` clause makes `net` a conditional-Normal kernel rather than a tensor: applied to an input it returns a distribution over `Target`, parameterised by a mean and a log-scale. Those two numbers are what the parameter network emits, so the concrete module behind `net` is -place [`MatrixNormal`](../api/continuous/families.md#quivers.continuous.families.ConditionalMatrixNormal) priors on each per-layer weight tensor. The two axes under `over (dom, cod)` bind positionally to the family's event axes: the input-side cardinality is the row axis and the output-side cardinality is the column axis, so the Kronecker covariance expresses independent row and column correlation in the weight matrix. - -The per-item input is itself a learnable morphism `X : Item -> H_in`; in a real workload `X` would be set from data via [`from_data`](../guides/dsl-overview.md). The composition `X >> W_1 >> W_2 >> W_3` is the full forward pass, materialising an `Item x H_out` score tensor under real-algebra matmul. +``` +Linear(1, 64) -> Tanh -> Linear(64, 64) -> Tanh -> Linear(64, 2) +``` -### Limitation +Its 4418 weights are the network's parameters. `hidden_dim` sets the width and `param_source` the architecture; `linear` is the default, and `identity` and `attention` are the other choices. Without `param_source=mlp` this kernel would map its input to the Normal's parameters through a single matrix, and the model would be a linear regression with a learned noise scale. -QVR's pure-composition surface under `composition real [level=algebra]` is strictly linear. There is no pointwise nonlinearity between composed weight matrices and no stochastic observation kernel on top of the discrete codomain `H_out`. The model expressed here is therefore a Bayesian *linear* network with matrix-normal priors on every layer's tensor. A deep nonlinear MLP with Bayesian weights is not currently expressible as a pure latent-morphism composition: a continuous-space surface would have to wrap each `W_l` in a continuous kernel carrying a nonlinearity inside its parameter network. The closest categorical form of the original BNN, with the matrix-normal priors actually entering the inference path, is the per-layer linear composition shown above. +`program bnn : Resp -> Resp` then does the only thing left. `observe y : Resp <- net(x)` applies the kernel to the per-row input and scores the observed response under the resulting Normal, accumulating over the `Resp` plate. `x` is a free variable: it never appears in a `sample` or `let`, so it is supplied as host data through the observations dict alongside `y`, exactly as a covariate would be in a regression. ## Try it -### SVI +> The step counts and NUTS budgets in the snippets below are illustrative: each block is sized to run in tens of seconds and demonstrate the API surface. Production fits typically need more steps, longer warmup, and multiple chains to converge. + +### Generating synthetic data ```python import torch from quivers.dsl import load torch.manual_seed(0) - prog = load("docs/examples/source/bnn.qvr") model = prog.morphism -# The model tensor materialises the Item x H_out score matrix. -# Fit it as a low-rank linear network to a target Y by gradient -# descent on the matmul output. N = 200 -H_out = 2 -Y = torch.randn(N, H_out) +x = torch.linspace(-3.0, 3.0, N).unsqueeze(-1) +y = torch.sin(3.0 * x) + 0.1 * torch.randn(N, 1) + +x_in = torch.zeros(N, 1) +observations = {"y": y, "x": x} +``` + +The target is a sine wave in noise, which is the point: it is a function no linear model can represent. `y` carries the `Target` event axis, so it is shaped `(N, 1)` rather than `(N,)`. The dummy `x_in` satisfies the program's `Resp -> Resp` signature; the real covariate enters through `observations["x"]`. -opt = torch.optim.Adam(prog.parameters(), lr=2e-2) -for _ in range(300): - opt.zero_grad() - loss = (model.tensor - Y).pow(2).mean() +### Fitting the network + +```python +optim = torch.optim.Adam(model.parameters(), lr=5e-3) + +for step in range(2000): + optim.zero_grad() + loss = -model.log_joint(x_in, observations)[0] loss.backward() - opt.step() + optim.step() -print("residual MSE:", (model.tensor - Y).pow(2).mean().item()) +print(f"final negative log-likelihood: {float(loss.detach()):.1f}") ``` -### NUTS posterior +[`log_joint`](../api/continuous/programs.md#quivers.continuous.programs.MonadicProgram.log_joint) returns the total log-joint broadcast across the batch axis, so every entry is the same number and `[0]` reads it. Summing instead would count the same scalar once per row. + +### Does the nonlinearity earn its keep? -Full Bayesian inference uses [`NUTSKernel`](../api/inference/mcmc.md#quivers.inference.mcmc.NUTSKernel) over the same model. For models declaring explicit `sample` priors NUTS samples them directly; models whose latents are `[role=latent]` parameters are lifted into a Normal-prior Bayesian model with [`bayesian_lift_parameters`](../api/inference/lifts.md#quivers.inference.lifts.bayesian_lift_parameters) so the standard MCMC machinery applies uniformly. +The honest comparison is against the best a linear model can do on the same data, scored the same way. Ordinary least squares gives the optimal line, and its residual spread gives the matching Gaussian noise level. ```python -import torch -from quivers.dsl import load -from quivers.inference import MCMC, NUTSKernel +design = torch.cat([torch.ones(N, 1), x], dim=-1) +beta = torch.linalg.lstsq(design, y).solution +resid = y - design @ beta +linear_nll = -torch.distributions.Normal(design @ beta, resid.std()).log_prob(y).sum() + +print(f"MLP negative log-likelihood: {float(loss.detach()):.1f}") +print(f"linear negative log-likelihood: {float(linear_nll):.1f}") +print(f"residual sd of the best line: {float(resid.std()):.3f}") +``` -torch.manual_seed(0) -prog = load("docs/examples/source/bnn.qvr") -model = prog.morphism +The line's residual spread lands near 0.7 against a true observation noise of 0.1, which is the quantitative form of the obvious: a straight line through a full period of a sine explains almost nothing, and the noise it infers is really the signal it cannot represent. The MLP's likelihood is several hundred nats better, and the gap is entirely attributable to the tanh layers, since the two models see identical data and differ only in the map from `x` to the Normal's parameters. -# Construct ``x`` and ``observations`` exactly as in the SVI block -# above. For models with no explicit ``sample`` priors, lift the -# parameters into a Bayesian model under unit Normal priors: -# from quivers.inference import bayesian_lift_parameters -# model, x, observations = bayesian_lift_parameters( -# model, x, observations, prior_scale=1.0, -# ) +### NUTS posterior over the weights -kernel = NUTSKernel(step_size=0.05, max_tree_depth=4, target_accept=0.8) -mc = MCMC(kernel, num_warmup=30, num_samples=30, num_chains=2) -result = mc.run(model, x, observations) +Nothing so far is Bayesian: the fit above is maximum likelihood over the weights. `bnn` declares no `sample` sites, so there is no prior in the source to sample. [`lift_from_log_prob`](../api/inference/lifts.md#quivers.inference.lifts.lift_from_log_prob) supplies one, putting a Normal prior on every weight and scoring the data through a log-density function of your choosing. The result is a [`MonadicProgram`](../api/continuous/programs.md#quivers.continuous.programs.MonadicProgram) whose sites are the weights themselves, which is exactly the Bayesian neural network the overview describes. -print("acceptance:", float(result.acceptance_rates.mean())) -print("divergences:", int(result.divergence_counts.sum())) -``` +```python +from quivers.inference import MCMC, NUTSKernel, lift_from_log_prob +def log_prob_fn(x_unused, y_obs): + return model.log_joint(x_in, {"y": y_obs, "x": x})[0].reshape(1) -## Categorical Perspective +lifted, lift_x, lift_obs = lift_from_log_prob( + model, + log_prob_fn=log_prob_fn, + parameter_prior_scale=1.0, + target_key="y", + observations={"y": y}, +) -Each weight tensor is a morphism in the discrete-object category whose prior measure on the hom-object $\mathbf{Kern}(\mathsf{H}_l, \mathsf{H}_{l + 1})$ is the [`MatrixNormal`](../api/continuous/families.md#quivers.continuous.families.ConditionalMatrixNormal) distribution. The forward pass is the composition +kernel = NUTSKernel(step_size=0.01, max_tree_depth=3, target_accept=0.8) +mc = MCMC(kernel, num_warmup=10, num_samples=10, num_chains=1) +result = mc.run(lifted, lift_x, lift_obs) -$$ -\mathsf{Item} \xrightarrow{X \mathbin{>>} W_1 \mathbin{>>} W_2 \mathbin{>>} W_3} \mathsf{H}_\text{out} -$$ +print(f"acceptance: {float(result.acceptance_rates.mean()):.2f}") +print(f"divergences: {int(result.divergence_counts.sum())}") +``` -in the real algebra. SVI's mean-field variational guide places an independent Normal posterior on every weight; predictive uncertainty is the marginal over weight samples drawn from this posterior. +`log_prob_fn` runs inside the lifted program's score step, after the sampled weights are substituted into the network's parameter slots, so it reads the current draw rather than the values the fit above left behind. Sampling 4418 weights is a genuine posterior over networks and the budget here is far too small to converge; it demonstrates the surface, not a usable posterior. -## See Also +## Categorical Perspective -- [Bayesian Linear Regression](bayesian-regression.md) for the single-layer linear special case. -- [DSL Guide](../guides/dsl-overview.md) for the morphism-valued prior surface and stochastic-kernel composition. +`net : Feature -> Target` is a Kleisli arrow $\mathbb{R} \to \mathcal{G}(\mathbb{R})$ in the [Giry monad](https://doi.org/10.1007/BFb0092872), sending an input to a Gaussian measure over the response. The MLP is not part of that arrow's type: it is the map from the domain into the family's parameter space, and composing it with the Gaussian's parameterisation is what makes the kernel's mean and scale depend nonlinearly on the input. +This is the difference between a nonlinearity in the parameter network and a nonlinearity between composed morphisms. Under a strictly linear algebra a chain `W_1 >> W_2 >> W_3` of tensor morphisms collapses: the composite is a single linear map, and the intermediate objects buy nothing but a rank bound. Placing the nonlinearity inside the kernel's parameter source avoids that collapse without leaving the V-Cat surface, because the kernel was never required to be linear in its input. The [`ParamSource`](../api/continuous/param_source.md#quivers.continuous.param_source.ParamSource) abstraction is the seam: `mlp`, `linear`, `attention`, and `identity` all present the same interface to the family and differ only in what they compute. -## References +Putting a prior on the weights is then a second, independent move. It lifts the deterministic parameter $\theta$ into a sample site, so the model becomes a mixture of networks $\int p(y \mid x, \theta)\, p(\theta)\, d\theta$ rather than a single one. + +## See Also -- David J. C. MacKay. 1992. The evidence framework applied to classification networks. *Neural Computation*, 4(3):448–472. +- [Deep Markov Model](deep-markov.md) for the same conditional-Normal kernels threaded through `scan` as a nonlinear state-space model. +- [Bayesian Regression](bayesian-regression.md) for the linear counterpart, where the coefficients carry declared priors in the source. diff --git a/docs/examples/continuous-hmm.md b/docs/examples/continuous-hmm.md index 601736c5..9656df2f 100644 --- a/docs/examples/continuous-hmm.md +++ b/docs/examples/continuous-hmm.md @@ -10,8 +10,8 @@ A continuous state-space model extends the HMM to continuous latent states and o object State : Real 16 object Obs : Real 8 -morphism transition : State -> State [scale=0.1] ~ Normal -morphism emission : State -> Obs [scale=0.1] ~ Normal +morphism transition : State -> State ~ Normal +morphism emission : State -> Obs ~ Normal program generative_step : State -> State sample s_new <- transition @@ -19,11 +19,11 @@ program generative_step : State -> State observe o <- emission(s_new) return s_new -morphism inference_cell : Obs * State -> State [scale=0.1] ~ Normal +morphism inference_cell : Obs * State -> State ~ Normal define filter = scan(inference_cell) -morphism decoder : State -> Obs [scale=0.1] ~ Normal +morphism decoder : State -> Obs ~ Normal define filter_and_reconstruct = scan(inference_cell) >> decoder @@ -35,13 +35,13 @@ export filter_and_reconstruct `object State : Real 16` and `object Obs : Real 8` introduce the two Euclidean spaces: a 16-dimensional latent and an 8-dimensional observation. -`morphism transition : State -> State [scale=0.1] ~ Normal` evolves the latent state by one time step under a Normal kernel whose mean is a learned linear function of the previous state and whose prior scale is 0.1. `morphism emission : State -> Obs [scale=0.1] ~ Normal` projects a state to an observation under the same kernel family. +`morphism transition : State -> State ~ Normal` evolves the latent state by one time step under a Normal kernel whose mean and scale are both produced from the previous state by the kernel's [`ParamSource`](../api/continuous/param_source.md#quivers.continuous.param_source.ParamSource). `morphism emission : State -> Obs ~ Normal` projects a state to an observation under the same kernel family. `program generative_step : State -> State` is a one-step monadic program: `sample s_new <- transition` draws the new latent state from the transition kernel, `observe o <- emission(s_new)` scores an observation against the emission kernel, and `return s_new` projects the program's joint kernel onto the new state. To unroll over time, this single-step program is composed with itself via `repeat` or threaded through `scan`. -`morphism inference_cell : Obs * State -> State [scale=0.1] ~ Normal` is a recurrent cell that incorporates a new observation into the running state estimate. `define filter = scan(inference_cell)` constructs a temporal-recurrence morphism that threads state across a sequence of observations. +`morphism inference_cell : Obs * State -> State ~ Normal` is a recurrent cell that incorporates a new observation into the running state estimate. `define filter = scan(inference_cell)` constructs a temporal-recurrence morphism that threads state across a sequence of observations. -`morphism decoder : State -> Obs [scale=0.1] ~ Normal` decodes a state back to observation space; `define filter_and_reconstruct = scan(inference_cell) >> decoder` composes the scan with the decoder so the exported pipeline filters and reconstructs in one composite. +`morphism decoder : State -> Obs ~ Normal` decodes a state back to observation space; `define filter_and_reconstruct = scan(inference_cell) >> decoder` composes the scan with the decoder so the exported pipeline filters and reconstructs in one composite. ## Try it diff --git a/docs/examples/deep-markov.md b/docs/examples/deep-markov.md index 6fcd6ab7..ff620444 100644 --- a/docs/examples/deep-markov.md +++ b/docs/examples/deep-markov.md @@ -21,11 +21,11 @@ object Hidden : Real 32 object State : Real 8 object Obs : Real 4 -morphism trans_mlp_1 : Driver * State -> Hidden [scale=0.5] ~ Normal -morphism trans_mlp_2 : Hidden -> State [scale=0.1] ~ Normal -morphism emit_mlp_1 : State -> Hidden [scale=0.5] ~ Normal -morphism emit_mlp_2 : Hidden -> Obs [scale=0.1] ~ Normal -morphism infer_cell : Obs * State -> State [scale=0.1] ~ Normal +morphism trans_mlp_1 : Driver * State -> Hidden [param_source=mlp] ~ Normal +morphism trans_mlp_2 : Hidden -> State [param_source=mlp] ~ Normal +morphism emit_mlp_1 : State -> Hidden [param_source=mlp] ~ Normal +morphism emit_mlp_2 : Hidden -> Obs [param_source=mlp] ~ Normal +morphism infer_cell : Obs * State -> State [param_source=mlp] ~ Normal define transition_cell = trans_mlp_1 >> trans_mlp_2 define emission = emit_mlp_1 >> emit_mlp_2 diff --git a/docs/examples/index.md b/docs/examples/index.md index 0e6ed304..ae570d65 100644 --- a/docs/examples/index.md +++ b/docs/examples/index.md @@ -27,7 +27,7 @@ All source files live under `docs/examples/source/`. - [Latent Dirichlet Allocation](lda.md): topic model with Dirichlet priors on per-document and per-topic distributions. - [Gaussian Mixture Model](mixture-model.md): finite mixture with grouped marginalization over the cluster label. - [Variational Autoencoder](vae.md): amortized inference over a continuous latent with neural decoder. -- [Bayesian Neural Network](bnn.md): feed-forward classifier with Normal priors over every weight. +- [Bayesian Neural Network](bnn.md): nonlinear regression through an MLP-parameterised Normal kernel, with Normal priors lifted over every weight. - [Parametric Partial Pooling](parametric-pooling.md): random effects from a parametric program template, with a labeled return tuple, a score-step sum-to-zero factor, and export selection. - [Probabilistic Matrix Factorization](pmf.md): low-rank Bayesian completion of a sparse rating matrix. - [Bilinear Tensor Contraction](tensor-contraction.md): neural-tensor-layer scoring of predicate-argument pairs via an operadic three-way contraction. diff --git a/docs/examples/lda.md b/docs/examples/lda.md index c04d7a83..2f1d506e 100644 --- a/docs/examples/lda.md +++ b/docs/examples/lda.md @@ -130,7 +130,7 @@ The discrete per-word topic $z : \mathsf{Topic}$ is integrated out by the [pushf ## See Also -- [Bayesian Gaussian Mixture Model](mixture-model.md) for a simpler grouped `marginalize` over a discrete latent. +- [Bayesian Gaussian Mixture Model](mixture-model.md) for a per-row mixture whose component assignment is integrated out in closed form by the likelihood, carrying no discrete latent and no `marginalize` block. ## References diff --git a/docs/examples/linear-gaussian-ssm.md b/docs/examples/linear-gaussian-ssm.md index ef957570..1a9fa20e 100644 --- a/docs/examples/linear-gaussian-ssm.md +++ b/docs/examples/linear-gaussian-ssm.md @@ -20,9 +20,9 @@ object Driver : Real 2 object State : Real 4 object Obs : Real 2 -morphism transition_cell : Driver * State -> State [scale=0.1] ~ Normal -morphism emission : State -> Obs [scale=0.1] ~ Normal -morphism filter_cell : Obs * State -> State [scale=0.1] ~ Normal +morphism transition_cell : Driver * State -> State ~ Normal +morphism emission : State -> Obs ~ Normal +morphism filter_cell : Obs * State -> State ~ Normal define generate = scan(transition_cell) >> emission define filter = scan(filter_cell) diff --git a/docs/examples/lstm-lm.md b/docs/examples/lstm-lm.md index 4e762426..b891d995 100644 --- a/docs/examples/lstm-lm.md +++ b/docs/examples/lstm-lm.md @@ -13,7 +13,7 @@ object Hidden : Real 128 morphism tok_embed : Token -> Embedded [role=embed] morphism gate_i, gate_f, gate_o : Embedded * Hidden -> Hidden ~ LogitNormal -morphism cell_cand : Embedded * Hidden -> Hidden [scale=0.5] ~ Normal +morphism cell_cand : Embedded * Hidden -> Hidden ~ Normal morphism lm_head : Hidden -> Token ~ Categorical program lstm_cell(x_t, h_prev) : Embedded * Hidden -> Hidden diff --git a/docs/examples/mixture-model.md b/docs/examples/mixture-model.md index e4835f98..ae6daba7 100644 --- a/docs/examples/mixture-model.md +++ b/docs/examples/mixture-model.md @@ -2,29 +2,30 @@ ## Overview -A finite [Gaussian mixture model](https://en.wikipedia.org/wiki/Mixture_model) assigns each observation to one of $K$ [Gaussian](https://en.wikipedia.org/wiki/Normal_distribution) components, with the per-row component drawn from a [Dirichlet](https://en.wikipedia.org/wiki/Dirichlet_distribution)-distributed mixing prior. This example demonstrates the canonical quivers idiom for finite mixtures: per-component means and scales are continuous latents on the `Component` plate, and the discrete per-row component assignment is integrated out by a scoped `marginalize` block whose body genuinely depends on the marginalized variable, yielding the canonical [log-sum-exp](https://en.wikipedia.org/wiki/LogSumExp) over $K$ classes at every observation: +A finite [Gaussian mixture model](https://en.wikipedia.org/wiki/Mixture_model) treats each observed value as a draw from one of $K$ [Gaussian](https://en.wikipedia.org/wiki/Normal_distribution) components. The components share three per-component vector parameters: the mixing weights on the simplex, the component locations, and the component scales. A [Dirichlet](https://en.wikipedia.org/wiki/Dirichlet_distribution) prior governs the mixing weights, and independent priors govern the locations and scales. + +Rather than sample a discrete per-row component label and integrate it out by hand, this example scores each row directly against the mixture through the [`MixtureNormal`](../api/continuous/families.md#quivers.continuous.families.ConditionalMixtureNormal) likelihood. The per-row marginal is the closed-form convex combination $$ -p(r_n) \;=\; \sum_{k=1}^{K} \mathrm{probs}[k] \; \mathcal{N}\!\bigl(r_n;\, \mu[k],\, \sigma[k]\bigr). +p(r_n) \;=\; \sum_{k=1}^{K} \mathrm{probs}[k] \; \mathcal{N}\!\bigl(r_n;\, \mu[k],\, \sigma[k]\bigr), $$ +so the model carries no discrete latent: `MixtureNormal` integrates the component assignment out analytically, evaluating the marginal by [log-sum-exp](https://en.wikipedia.org/wiki/LogSumExp) over the $K$ components. + ## QVR Source ```qvr composition log_prob [level=algebra] object Component : FinSet 3 -object Item : FinSet 8 object Resp : FinSet 100 program gmm(alpha : Real) : Resp -> Resp sample probs <- Dirichlet(alpha) [over=Component] sample mu : Component <- Normal(0.0, 5.0) sample sigma : Component <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) - marginalize cls : Component <- Categorical(probs) [over=Item, reduction=logsumexp] - observe r : Resp <- Normal(mu[cls], sigma[cls]) [via=idx] + observe r : Resp <- MixtureNormal(probs, mu, sigma) return probs @@ -33,19 +34,15 @@ export gmm ## Walkthrough -`composition log_prob [level=algebra]` selects the log-probability semiring so the program's `Score` effect accumulates log-densities additively. `object Component : FinSet 3`, `object Item : FinSet 8`, `object Resp : FinSet 100` declare the three discrete plates: $K = 3$ components, $I = 8$ item groups, $N = 100$ observed rows. `program gmm(alpha : Real) : Resp -> Resp` parameterises the program by the Dirichlet concentration. +`composition log_prob [level=algebra]` selects the log-probability semiring so the program's `Score` effect accumulates log-densities additively. `object Component : FinSet 3` and `object Resp : FinSet 100` declare the two discrete plates: $K = 3$ mixture components and $N = 100$ observed rows. `program gmm(alpha : Real) : Resp -> Resp` parameterises the program by the Dirichlet concentration. -`sample probs <- Dirichlet(alpha) [over=Component]` draws the mixing weights as a single point on the `Component` simplex; `over=Component` names the family's event axis. `sample mu : Component <- Normal(0.0, 5.0)` and `sample sigma : Component <- HalfNormal(1.0)` draw the per-component mean and scale as plate-bound continuous latents. `sample idx : Resp <- HalfNormal(1.0)` registers the per-row fibration site that names the runtime map from `Resp` into the `Item` grouping plate. +The three `sample` steps draw the shared per-component parameters: -The scoped `marginalize` block +- `sample probs <- Dirichlet(alpha) [over=Component]` draws the mixing weights as a single point on the `Component` simplex; `over=Component` names the family's event axis (Dirichlet event-rank 1). +- `sample mu : Component <- Normal(0.0, 5.0)` draws the $K$ component locations as plate-bound continuous latents, one per component. +- `sample sigma : Component <- HalfNormal(1.0)` draws the $K$ positive component scales the same way. - -```qvr -marginalize cls : Component <- Categorical(probs) [over=Item, reduction=logsumexp] - observe r : Resp <- Normal(mu[cls], sigma[cls]) [via=idx] -``` - -introduces the per-row component latent `cls : Component` under a [Categorical](https://en.wikipedia.org/wiki/Categorical_distribution) prior parameterised by `probs`. The body's `Normal(mu[cls], sigma[cls])` looks up the chosen component's mean and scale and scores the observed row `r` against that per-class Gaussian. The `[over=Item]` grouping plate accumulates each observation into its `Item` bucket; `[via=idx]` names the runtime fibration from each row into its item group. The `[reduction=logsumexp]` reduction integrates `cls` out by pushforward along the projection $\Phi \times \mathsf{Component} \to \Phi$, recovering the closed-form per-row mixture marginal $\sum_k \mathrm{probs}[k]\,\mathcal{N}(r_n;\,\mu[k],\,\sigma[k])$. At the end of the scope `cls` falls out of scope. +`observe r : Resp <- MixtureNormal(probs, mu, sigma)` scores each observed row against the $K$-component mixture the three shared vectors define. `MixtureNormal` takes the mixing weights, the locations, and the scales as three per-component vectors, each broadcast across every row of the `Resp` plate, and returns the per-row marginal $\sum_k \mathrm{probs}[k]\,\mathcal{N}(r_n;\,\mu[k],\,\sigma[k])$ in closed form. No component-assignment latent is sampled: the mixture likelihood integrates it out. `return probs` projects the program's joint kernel onto the mixing-weight site. @@ -72,13 +69,12 @@ true_sigma = torch.tensor([0.5, 0.7, 0.4]) comps = torch.distributions.Categorical(true_probs).sample((N,)) r = torch.distributions.Normal(true_mu[comps], true_sigma[comps]).sample() -idx = torch.randint(0, 8, (N,)) -observations = {"r": r, "idx": idx, "probs": true_probs} +observations = {"r": r, "probs": true_probs} x_in = torch.zeros(N, 1) ``` -The simplex-supported `probs` site is supplied via `observations` so the grouped marginalize block sees a `(K,)` mixing prior at every batch position. The Gaussian per-component parameters `mu` and `sigma` remain unobserved and are recovered by SVI. +The synthetic rows are drawn by sampling a component per row and then a value from that component's Gaussian, but the model never sees the component labels: only the rows `r` and the fixed mixing weights `probs` enter `observations`. The per-component locations `mu` and scales `sigma` remain unobserved latents and are recovered by SVI. ### SVI fit @@ -86,19 +82,22 @@ The simplex-supported `probs` site is supplied via `observations` so the grouped from quivers.inference import AutoNormalGuide, ELBO, SVI torch.manual_seed(1) -guide = AutoNormalGuide( - model, observed_names={"r", "idx", "probs"}, -) +guide = AutoNormalGuide(model, observed_names={"r", "probs"}) optim = torch.optim.Adam( - list(model.parameters()) + list(guide.parameters()), lr=5e-2, + list(model.parameters()) + list(guide.parameters()), lr=1e-1, ) svi = SVI(model, guide, optim, ELBO(num_particles=1)) -losses = [svi.step(x_in, observations) for _ in range(300)] -print(f"initial loss: {losses[0]:.2f}") -print(f"final loss: {losses[-1]:.2f}") +losses = [svi.step(x_in, observations) for _ in range(1200)] +recovered = sorted(guide.loc_mu.detach().flatten().tolist()) +print(f"initial loss: {losses[0]:.2f}") +print(f"final loss: {losses[-1]:.2f}") +print("recovered means: [" + ", ".join(f"{m:.2f}" for m in recovered) + "]") +print("true means: [-3.00, 0.00, 3.00]") ``` +The variational locations `guide.loc_mu` hold the recovered component means. Because the mixture ELBO is multimodal, the fit is sensitive to initialisation and to the mixing weights held fixed at `probs`; a sharper separation between components makes the locations easier to recover. + ### NUTS posterior ```python @@ -116,8 +115,8 @@ print(f"divergences: {int(result.divergence_counts.sum())}") ## Categorical Perspective -The discrete latent `cls : Component` is integrated out by [pushforward](https://en.wikipedia.org/wiki/Pushforward_measure) along the projection $\Phi \times \mathsf{Component} \to \Phi$. The grouped marginalize block is the [right Kan extension](https://ncatlab.org/nlab/show/Kan+extension) of the per-class log-likelihood along the per-row fibration $\mathsf{Resp} \to \mathsf{Item}$ in $\mathbf{Kern}$, followed by a [log-sum-exp](https://en.wikipedia.org/wiki/LogSumExp) reduction along the `Component` axis weighted by the categorical prior implied by the Dirichlet. +The per-row likelihood `MixtureNormal(probs, mu, sigma)` is a [Kleisli arrow](https://en.wikipedia.org/wiki/Kleisli_category) $\mathsf{Resp} \to \mathsf{Resp}$ in the [Giry monad](https://doi.org/10.1007/BFb0092872). It is a finite convex combination of the $K$ Gaussian component measures, weighted by the categorical measure $\mathrm{probs}$ on `Component`: the Giry-monad mixture operation that draws a component from $\mathrm{Categorical}(\mathrm{probs})$, then a value from the chosen Gaussian, and keeps the marginal on the value. Equivalently, `MixtureNormal` is the [pushforward](https://en.wikipedia.org/wiki/Pushforward_measure) of the joint component-and-value measure along the projection $\mathsf{Component} \times \mathbb{R} \to \mathbb{R}$, which is exactly the closed-form marginal $\sum_k \mathrm{probs}[k]\,\mathcal{N}(\cdot;\,\mu[k],\,\sigma[k])$. Because the component index is integrated out inside the likelihood rather than sampled, the program carries no discrete latent to marginalise and every site it holds is continuous. ## See Also -- [Latent Dirichlet Allocation](lda.md), the topic-model generalization. +- [Latent Dirichlet Allocation](lda.md), the grouped discrete-mixture generalisation whose per-word topic assignment is integrated out by a scoped `marginalize` block. diff --git a/docs/examples/montague-nli.md b/docs/examples/montague-nli.md index 952bf2c1..56407a39 100644 --- a/docs/examples/montague-nli.md +++ b/docs/examples/montague-nli.md @@ -2,62 +2,79 @@ ## Overview -A two-stage natural-language-inference architecture composed entirely out of QVR's weighted-deduction surface: a Montague-style grammar that derives logical forms from token spans, then an entailment prover that closes the LFs under modus ponens and functorial substitution. Everything is declared in `atoms` and `binders` plus `rule` sequents; no grammar formalism or proof system is baked into the language. +A two-stage natural-language-inference architecture composed entirely out of QVR's weighted-deduction surface: a Montague-style grammar that derives a logical form for every token span, then an entailment prover that closes those logical forms under three syllogistic rules. Everything is declared in `atoms` and `binders` plus `rule` sequents; no grammar formalism and no proof system is baked into the language. + +The two halves share one term language, so the prover's items are literally the terms the grammar builds. That is what lets a single Adam step move the lexicon's log-weights in response to an entailment error. ## QVR Source ```qvr object Term : FinSet 8 -deduction Montague : Term -> Term [semiring=LogProb, start=S, depth=10] - atoms NP, S, N, Fwd, Bwd, span, App, Var, dog_p, cat_p, animal_p, bark_p, walk_p, forall_t, exists_t, implies_t, and_t +deduction Montague : Term -> Term [semiring=LogProb, start=S, depth=12] + atoms S, N, VP, Nom, Art, Cop, DetEvery, DetSome, QEvery, QSome, span, App, Var, Every, Some, dog_p, cat_p, animal_p, bark_p, walk_p, every_q, some_q, an_q, is_q binders Lam - rule fwd_app : span(I, K, Fwd(X, Y), F), span(K, J, X, A) |- span(I, J, Y, App(F, A)) #[learnable] - rule bwd_app : span(I, K, X, A), span(K, J, Bwd(X, Y), F) |- span(I, J, Y, App(F, A)) #[learnable] + rule every_np : span(I, K, DetEvery, D), span(K, J, N, P) |- span(I, J, QEvery, P) #[learnable] + rule some_np : span(I, K, DetSome, D), span(K, J, N, P) |- span(I, J, QSome, P) #[learnable] + rule every_s : span(I, K, QEvery, P), span(K, J, VP, Q) |- span(I, J, S, Every(P, Q)) #[learnable] + rule some_s : span(I, K, QSome, P), span(K, J, VP, Q) |- span(I, J, S, Some(P, Q)) #[learnable] + rule art_n : span(I, K, Art, D), span(K, J, N, P) |- span(I, J, Nom, P) #[learnable] + rule cop_nom : span(I, K, Cop, D), span(K, J, Nom, P) |- span(I, J, VP, P) #[learnable] lexicon - "dog" : N = Lam(x, App(dog_p, Var(x))) #[learnable] - "cat" : N = Lam(x, App(cat_p, Var(x))) #[learnable] - "animal" : N = Lam(x, App(animal_p, Var(x))) #[learnable] - "barks" : Bwd(NP, S) = Lam(x, App(bark_p, Var(x))) #[learnable] - "walks" : Bwd(NP, S) = Lam(x, App(walk_p, Var(x))) #[learnable] - "every" : Fwd(N, Fwd(Bwd(NP, S), S)) = Lam(P, Lam(Q, App(forall_t, Lam(x, App(App(implies_t, App(Var(P), Var(x))), App(Var(Q), Var(x))))))) #[learnable] - "some" : Fwd(N, Fwd(Bwd(NP, S), S)) = Lam(P, Lam(Q, App(exists_t, Lam(x, App(App(and_t, App(Var(P), Var(x))), App(Var(Q), Var(x))))))) #[learnable] - -deduction Prover : Term -> Term [semiring=LogProb, depth=8] - atoms Claim, Implies, App, Var + "dog" : N = Lam(x, App(dog_p, Var(x))) #[learnable] + "cat" : N = Lam(x, App(cat_p, Var(x))) #[learnable] + "animal" : N = Lam(x, App(animal_p, Var(x))) #[learnable] + "barks" : VP = Lam(x, App(bark_p, Var(x))) #[learnable] + "walks" : VP = Lam(x, App(walk_p, Var(x))) #[learnable] + "every" : DetEvery = every_q #[learnable] + "some" : DetSome = some_q #[learnable] + "an" : Art = an_q #[learnable] + "is" : Cop = is_q #[learnable] + +deduction Prover : Term -> Term [semiring=LogProb, depth=12] + atoms Claim, Every, Some, Nonempty, App, Var binders Lam - rule modus_ponens : Claim(P), Claim(Implies(P, Q)) |- Claim(Q) #[learnable] - rule app_subst : Claim(App(F, X)), Claim(Implies(X, Y)) |- Claim(App(F, Y)) #[learnable] + rule barbara : Claim(Every(P, Q)), Claim(Every(Q, R)) |- Claim(Every(P, R)) #[learnable] + rule darii : Claim(Some(P, Q)), Claim(Every(Q, R)) |- Claim(Some(P, R)) #[learnable] + rule ex_import : Claim(Every(P, Q)), Claim(Nonempty(P)) |- Claim(Some(P, Q)) #[learnable] program fit_grammar : Term -> Term let chart = parse(Montague, sentence) score log_Z = chart.goal_weight() return log_Z -program fit_pipeline : Term -> Term - let pipeline = compose(Montague, Prover) - let chart = parse(pipeline, sentence) - score log_Z = chart.goal_weight() - return log_Z - export fit_grammar ``` ## Walkthrough -The grammar half of the module declares atomic category constructors (`NP`, `S`, `N`), slash constructors (`Fwd`, `Bwd`), the chart-item constructor `span(I, J, X, F)` that packages a derivation covering tokens `[I, J)` of category `X` with logical form `F`, a function-application LF combinator `App`, a variable-occurrence constructor `Var`, the predicate constants `dog_p`, `cat_p`, `animal_p`, `bark_p`, `walk_p`, and the logical connectives `forall_t`, `exists_t`, `implies_t`, `and_t` used in determiner denotations. The `binders Lam` block tells the compiler that `Lam`'s first argument is a binding site, so every bound variable is alpha-renamed to a fresh canonical symbol per term construction and structural equality on the chart is [alpha-equivalence](https://en.wikipedia.org/wiki/Lambda_calculus#Alpha_equivalence) on the surface. +The grammar half declares the categories (`S`, `N`, `VP`, `Nom`, and the closed-class categories `Art`, `Cop`, `DetEvery`, `DetSome`), the chart-item constructor `span(I, J, X, F)` that packages a derivation covering tokens `[I, J)` of category `X` with logical form `F`, the logical-form constructors `App`, `Var`, `Every`, and `Some`, and the predicate constants `dog_p`, `cat_p`, `animal_p`, `bark_p`, `walk_p`. The `binders Lam` block tells the compiler that `Lam`'s first argument is a binding site, so every bound variable is alpha-renamed to a fresh canonical symbol per lexicon entry and structural equality on the chart is [alpha-equivalence](https://en.wikipedia.org/wiki/Lambda_calculus#Alpha_equivalence) on the surface. This matters for the prover: because "dog" is compiled once, the predicate it denotes is one canonical term, and the prover's pattern variables bind it by structural equality wherever it occurs. + +### Why the sentence LF is built in normal form + +One constraint on the surface drives the whole design, so it is worth stating precisely. A deduction rule's conclusion is a constructor tree over the rule's pattern variables, and the runtime instantiates it with [`instantiate`](../api/stochastic/agenda.md#quivers.stochastic.agenda.instantiate), which performs structural substitution and nothing else. There is no beta-rule anywhere: a rule conclusion cannot call a function, no chart item is normalised on the way in, and the `binders` machinery only alpha-renames lexicon logical forms at compile time. + +Thus a determiner cannot denote a continuation-form lambda term such as `Lam(P, Lam(Q, App(forall_t, ...)))` and rely on later beta-reduction against its arguments. Nothing would reduce the redex; the chart would carry `App(App(Lam(...), dog_LF), bark_LF)` forever, and no prover rule could pattern-match through the unreduced application to find the quantifier. -Forward and backward application (`fwd_app` and `bwd_app`) are the two sequents that combine a slash-typed span with its complement, building the conclusion's logical form by applying the head's LF to its argument's LF via `App`. The lexicon ships learnable log-weights per entry: common nouns and intransitive verbs are unary predicates `Lam(x, App(p, Var(x)))`, and determiners are honest [generalised quantifiers](https://en.wikipedia.org/wiki/Generalized_quantifier) in continuation form, with bound variables `P`, `Q`, `x` introduced by `Lam` and substituted by the application rules. See [Montague (1973)](https://doi.org/10.1007/978-94-010-2506-5_10) for the type-driven semantic compositionality this fragment instantiates. +The grammar therefore builds the sentence logical form in normal form directly, in the [generalized quantifier](https://en.wikipedia.org/wiki/Generalized_quantifier) style of [Barwise and Cooper (1981)](https://doi.org/10.1007/BF00350139): a determiner denotes a relation between two sets, written here as the binary constructors `Every(P, Q)` and `Some(P, Q)`. The determiner's quantificational force rides on its category, and the sentence rule that consumes that category (`every_s` or `some_s`) builds the matching constructor. The restrictor and the scope remain genuine lambda terms, so the binder machinery does real work at the predicate level, where it is what makes `Lam(x, App(dog_p, Var(x)))` a single canonical object rather than a name-dependent one. See [Montague (1973)](https://doi.org/10.1007/978-94-010-2506-5_10) for the type-driven compositionality this fragment instantiates. -The prover half closes the resulting `Claim` items under modus ponens and a functorial-substitution rule: from `Claim(App(F, X))` and `Claim(Implies(X, Y))`, conclude `Claim(App(F, Y))`. Both deductions share the `binders Lam` declaration so capture-avoiding substitution lives at the compile-time canonical-variable layer. +Two consequences follow, and both are visible in the source. First, the closed-class words carry no logical form of their own: `every`, `some`, `an`, and `is` name themselves with a constant (`every_q` and friends) so that their log-weights stay learnable, while their category is what drives the rules. This is the same convention the [CCG example](ccg.md) uses. Second, the article and the copula are semantically vacuous, so `art_n` and `cop_nom` pass the noun's predicate through unchanged, which is what makes `every dog is an animal` and `every animal barks` come out as `Every(DOG, ANIMAL)` and `Every(ANIMAL, BARK)` over the *same* `ANIMAL` term. -The two top-level programs expose the grammar and the composed pipeline as differentiable parsers. `fit_grammar` calls `parse(Montague, sentence)` against the host-supplied token list and `score`s the chart's goal weight (the inside log-marginal) into the program. `fit_pipeline` uses `compose(Montague, Prover)` to feed the grammar's logical forms into the prover's chart, so the prover's goal weight inherits the grammar's gradient flow into the lexicon's log-weights. +### The prover + +The prover's items are `Claim(phi)` for `phi` a sentence logical form the grammar built, plus `Nonempty(P)` for an explicit existence premise. Read the constructors as + +$$\mathtt{Every}(P, Q) \;=\; \forall x.\, P(x) \to Q(x), \qquad \mathtt{Some}(P, Q) \;=\; \exists x.\, P(x) \wedge Q(x), \qquad \mathtt{Nonempty}(P) \;=\; \exists x.\, P(x).$$ + +All three rules are classically valid on that reading. `barbara` (every $P$ is $Q$, every $Q$ is $R$, thus every $P$ is $R$) needs no existence assumption. `darii` (some $P$ is $Q$, every $Q$ is $R$, thus some $P$ is $R$) is valid because the existential premise carries its own witness. + +`ex_import` is the one that needs care. `Every(P, Q)` alone does **not** entail `Some(P, Q)`: on the standard first-order reading a universal over an empty restrictor is vacuously true while the corresponding existential is false, so `every dog barks` does not by itself license `some dog barks`. The inference is valid only under [existential import](https://en.wikipedia.org/wiki/Square_of_opposition#Existential_import), that is, only when the restrictor is known to be non-empty. The rule therefore takes `Nonempty(P)` as a second premise rather than smuggling the assumption into the rule's shape, and a caller that wants the inference has to supply that premise. The Python below supplies it explicitly, for exactly one noun, and says so. ## Try it -The grammar denotes a sentence as a logical form: a structured term over the user-declared free term algebra. Lambda terms with bound variables (`Lam(x, App(bark_p, Var(x)))`) are honest object-level lambda calculus: the `binders Lam` declaration tells the compiler that `Lam`'s first argument is a binding site, so every bound variable is alpha-renamed to a fresh canonical symbol at compile time and the chart's structural identity collapses alpha-equivalent terms. +Every `#[learnable]` lexicon entry and every `#[learnable]` rule exposes a real `nn.Parameter` on the compiled [`DeductionSystem`](../api/stochastic/agenda.md#quivers.stochastic.agenda.DeductionSystem). The system is callable: `ded(sentence)` returns a [`ChartView`](../api/stochastic/agenda.md#quivers.stochastic.agenda.ChartView) whose [`goal_weight`](../api/stochastic/agenda.md#quivers.stochastic.agenda.ChartView.goal_weight) is the differentiable inside log-marginal at the start symbol. -### SVI +### MAP fit on the grammar ```python import torch @@ -67,142 +84,195 @@ from quivers.stochastic.deduction import adam_fit_deduction torch.manual_seed(0) prog = load("docs/examples/source/montague_nli.qvr") grammar = prog.deductions["Montague"] -prover = prog.deductions["Prover"] +prover = prog.deductions["Prover"] -# Training corpus: pairs of (premise, hypothesis) sentences whose -# joint NLI label we want the grammar to fit. The Montague chart -# derives a logical form for each sentence; the prover then closes -# the resulting Claims under modus ponens to check entailment. corpus = [ - (["every", "dog", "barks"], ["some", "dog", "barks"]), - (["every", "cat", "walks"], ["some", "cat", "walks"]), + ["every", "dog", "barks"], + ["some", "dog", "barks"], + ["every", "cat", "walks"], + ["every", "dog", "is", "an", "animal"], + ["some", "dog", "is", "an", "animal"], + ["every", "animal", "barks"], ] -# Touch every premise + hypothesis once so the deduction's lazy -# rule-weight ParameterDicts allocate every binding tuple they -# will see during fitting. -for premise, hypothesis in corpus: - grammar(premise).goal_weight() - grammar(hypothesis).goal_weight() - -# Adam over the grammar's lexicon + rule log-weights. The loss is -# minus the corpus log-marginal: a higher chart total at the start -# symbol of each sentence under the current parameters. -all_sentences = [s for pair in corpus for s in pair] +# Touch every sentence once so the deduction's lazily allocated +# lexicon- and rule-weight ParameterDicts materialise every binding +# tuple the fit will see. +for sentence in corpus: + grammar(sentence).goal_weight() + history = adam_fit_deduction( - grammar, all_sentences, steps=200, lr=5e-2, prior_scale=1.0, + grammar, corpus, steps=200, lr=5e-2, prior_scale=1.0, ) print(f"loss: {history[0]:.2f} -> {history[-1]:.2f}") -for sentence in all_sentences: - lf = next( - item[4] for item, _ in grammar(sentence).chart.items() - if isinstance(item, tuple) and item[:1] == ("span",) - and len(item) >= 5 and item[3] == ("atom", "S") - ) + +def sentence_lf(sentence): + """The logical form the grammar assigns to the whole sentence, + paired with the chart's goal weight for that derivation.""" + chart = grammar(sentence) + target = ("span", 0, len(sentence), ("atom", "S")) + for item, _ in chart.chart.items(): + if isinstance(item, tuple) and item[:4] == target: + return item[4], chart.goal_weight() + raise ValueError(f"no S-derivation for {' '.join(sentence)!r}") + + +for sentence in corpus: + lf, _ = sentence_lf(sentence) print(f"LF({' '.join(sentence)}) = {lf}") ``` -### Entailment via the Prover (full NLI loss) +Every logical form printed there is in normal form: its head is `Every` or `Some`, never an `App` whose function is a `Lam`. That is the property the prover depends on. -The prover deduction closes the Claims under modus ponens and a functorial substitution rule. An NLI label is the prover's goal weight at `Claim(hypothesis_lf)` after seeding the chart with `Claim(premise_lf)`. +### Entailment via the prover -```python -import torch -from quivers.dsl import load -from quivers.stochastic.deduction import adam_fit_deduction +The prover takes an axiom list of `(item, log_weight)` pairs. Seeding `Claim(premise_lf)` at the *grammar's* goal weight for that premise, rather than at a constant, is what puts the lexicon's parameters into the prover's chart: the weight of every derived claim is a sum that includes it, so a gradient at `Claim(hypothesis_lf)` reaches back through the grammar. -torch.manual_seed(0) -prog = load("docs/examples/source/montague_nli.qvr") -grammar = prog.deductions["Montague"] -prover = prog.deductions["Prover"] +The NLI corpus below is three entailments and one non-entailment. The third pair is the interesting one. It is licensed only by the explicit `Nonempty(DOG)` premise, which asserts that the dog predicate has a witness; without it, `every dog barks` would not entail `some dog barks`, and the prover would (correctly) derive nothing. -def derive_lf(sentence): - """Run the grammar; return the logical form at span(0, n, S).""" - chart = grammar(sentence) - n = len(sentence) - target = ("span", 0, n, ("atom", "S")) +```python +def noun_pred(noun): + """The one-place predicate a common noun denotes.""" + chart = grammar([noun]) + target = ("span", 0, 1, ("atom", "N")) for item, _ in chart.chart.items(): if isinstance(item, tuple) and item[:4] == target: - return item[4] if len(item) > 4 else None - return None - -def entailment_log_score(premise, hypothesis): - """Compute log P(prover derives Claim(hypothesis_lf) given - premise_lf as an axiom). The prover's goal weight at the - hypothesis Claim is the NLI score; gradients flow back into the - grammar's lexicon and rule weights.""" - p_lf = derive_lf(premise) - h_lf = derive_lf(hypothesis) - if p_lf is None or h_lf is None: - return torch.tensor(float("-inf")) - # Seed the prover with the premise as an axiom; the prover - # closes the chart and reports the weight at Claim(h_lf). - prover_chart = prover([(("Claim", p_lf), torch.tensor(0.0))]) - return prover_chart.try_weight(("Claim", h_lf)) - -# A toy entailment corpus: (premise, hypothesis, label) triples -# where label = 1 when the hypothesis is entailed. + return item[4] + raise ValueError(f"{noun!r} is not a common noun in the lexicon") + + +# (premises, nouns-asserted-non-empty, hypothesis, label) nli = [ - (["every", "dog", "barks"], ["some", "dog", "barks"], 1), - (["every", "cat", "walks"], ["some", "cat", "walks"], 1), - (["every", "dog", "barks"], ["every", "cat", "walks"], 0), + # Barbara. Valid with no existence assumption. + ([["every", "dog", "is", "an", "animal"], ["every", "animal", "barks"]], + [], ["every", "dog", "barks"], 1), + # Darii. The existential premise carries its own witness. + ([["some", "dog", "is", "an", "animal"], ["every", "animal", "barks"]], + [], ["some", "dog", "barks"], 1), + # Subalternation. Valid ONLY given the existence premise, which + # is why "dog" is listed as non-empty here and nowhere else. + ([["every", "dog", "barks"]], ["dog"], ["some", "dog", "barks"], 1), + # Not entailed: nothing in the premises is about cats. + ([["every", "dog", "is", "an", "animal"], ["every", "animal", "barks"]], + [], ["every", "cat", "walks"], 0), ] + +def entailment_score(premises, nonempty_nouns, hypothesis): + """Score Claim(hypothesis_lf) against the prover's chart. + + Returns the log-weight at the hypothesis claim together with the + log-normaliser over every claim the prover holds, so that their + difference is log p(hypothesis | premises): a proper conditional + distribution over the prover's conclusions. The log-weight is the + LogProb semiring's zero (-inf) when the prover derives no such + claim.""" + axioms = [] + for premise in premises: + lf, log_z = sentence_lf(premise) + axioms.append((("Claim", lf), log_z)) + for noun in nonempty_nouns: + axioms.append((("Claim", ("Nonempty", noun_pred(noun))), torch.zeros(()))) + h_lf, _ = sentence_lf(hypothesis) + chart = prover(axioms) + claims = torch.stack( + [w for item, w in chart.chart.items() if item[0] == "Claim"] + ) + return chart.try_weight(("Claim", h_lf)), torch.logsumexp(claims, dim=0) + + +# Touch every prover chart once so its rule-weight ParameterDicts +# allocate every binding tuple before the optimiser reads them. +for premises, nonempty, hypothesis, label in nli: + entailment_score(premises, nonempty, hypothesis) + params = list(grammar.parameters()) + list(prover.parameters()) -optim = torch.optim.Adam(params, lr=5e-2) +optim = torch.optim.Adam(params, lr=5e-2) + +history, grad_at_start = [], [] for step in range(150): optim.zero_grad() loss = torch.zeros(()) - for premise, hypothesis, label in nli: - score = entailment_log_score(premise, hypothesis) - if torch.isfinite(score): - # Bernoulli NLI loss: -log p(label | score). - log_p_entailed = score - log_p_not = torch.log1p(-torch.exp(score)).clamp(min=-50.0) \ - if score.item() < 0 else torch.tensor(-50.0) - loss = loss - (label * log_p_entailed + (1 - label) * log_p_not) + for premises, nonempty, hypothesis, label in nli: + score, log_norm = entailment_score(premises, nonempty, hypothesis) + if label == 1: + assert torch.isfinite(score), ( + f"{' '.join(hypothesis)!r} is labelled entailed but the " + f"prover cannot derive it" + ) + loss = loss - (score - log_norm) + else: + # The prover is sound on this fragment, so it derives no + # claim for a non-entailment at all: p(entailed) is + # exactly zero and the Bernoulli term -log(1 - 0) is + # exactly zero. Assert the soundness rather than add a + # term that is identically zero. + assert not torch.isfinite(score), ( + f"{' '.join(hypothesis)!r} is labelled not entailed but the " + f"prover derived it" + ) loss.backward() + if step == 0: + grad_at_start = [ + p for p in params if p.grad is not None and float(p.grad.abs().max()) > 0 + ] + history.append(float(loss.detach())) optim.step() -print("final loss:", float(loss.detach())) +print(f"NLI loss: {history[0]:.3f} -> {history[-1]:.3f}") +assert torch.isfinite(loss), "loss must be finite" +assert history[-1] < history[0], "the fit must reduce the NLI loss" + +# Gradients are read at the first step, before the fit drives them +# toward zero by converging. +print(f"parameters with a non-zero gradient: {len(grad_at_start)} / {len(params)}") +assert grad_at_start, "no parameter received a gradient" + +for premises, nonempty, hypothesis, label in nli: + score, log_norm = entailment_score(premises, nonempty, hypothesis) + p_entail = float(torch.exp(score - log_norm).detach()) + print(f"p(entailed | {' '.join(hypothesis)}) = {p_entail:.3f} (label {label})") ``` +The final loop reports a calibrated probability per pair. The non-entailment sits at exactly `0.000`, because the prover derives no claim for it at all rather than because it learned to score it low; the prover is sound on this fragment by construction, so the negative supplies no gradient. The three entailments are what the fit actually moves. + ### NUTS posterior -Full Bayesian inference uses [`NUTSKernel`](../api/inference/mcmc.md#quivers.inference.mcmc.NUTSKernel) over the same model. For models declaring explicit `sample` priors NUTS samples them directly; models whose latents are `[role=latent]` parameters are lifted into a Normal-prior Bayesian model with [`bayesian_lift_parameters`](../api/inference/lifts.md#quivers.inference.lifts.bayesian_lift_parameters) so the standard MCMC machinery applies uniformly. +Full Bayesian inference over the grammar's log-weights uses [`NUTSKernel`](../api/inference/mcmc.md#quivers.inference.mcmc.NUTSKernel). [`nuts_program_from_deduction`](../api/stochastic/deduction.md#quivers.stochastic.deduction.nuts_program_from_deduction) lifts every learnable parameter into a [`Normal`](../api/continuous/families.md) sample site and adds the corpus log-marginal to the joint via a `score` step, so the standard MCMC machinery applies unchanged. ```python -import torch -from quivers.dsl import load from quivers.inference import MCMC, NUTSKernel +from quivers.stochastic.deduction import nuts_program_from_deduction torch.manual_seed(0) -prog = load("docs/examples/source/montague_nli.qvr") -model = prog.morphism - -# Construct ``x`` and ``observations`` exactly as in the SVI block -# above. For models with no explicit ``sample`` priors, lift the -# parameters into a Bayesian model under unit Normal priors: -# from quivers.inference import bayesian_lift_parameters -# model, x, observations = bayesian_lift_parameters( -# model, x, observations, prior_scale=1.0, -# ) +model, x, observations = nuts_program_from_deduction( + grammar, corpus, prior_scale=1.0, +) kernel = NUTSKernel(step_size=0.05, max_tree_depth=4, target_accept=0.8) -mc = MCMC(kernel, num_warmup=30, num_samples=30, num_chains=2) +mc = MCMC(kernel, num_warmup=30, num_samples=30, num_chains=2) result = mc.run(model, x, observations) print("acceptance:", float(result.acceptance_rates.mean())) print("divergences:", int(result.divergence_counts.sum())) ``` - ## Categorical Perspective -Each `deduction` block denotes a weighted relation in the [agenda-based deduction semiring](../semantics/composition-rules.md): an arrow $\mathrm{Term} \to \mathrm{Term}$ in the [LogProb algebra](../semantics/algebras.md) whose underlying tensor is the chart of derivable items keyed by their derivation log-weights. Composing grammar with prover is composition in the same enriched category, so the gradient of the prover's goal weight flows back through the grammar's lexicon entries during training. +Each `deduction` block denotes a weighted relation in the [agenda-based deduction semiring](../semantics/composition-rules.md): an arrow $\mathrm{Term} \to \mathrm{Term}$ in the [LogProb algebra](../semantics/algebras.md) whose underlying tensor is the chart of derivable items keyed by their derivation log-weights. + +The two deductions are chained by hand rather than by a `compose(...)` step, and the reason is worth stating. `compose(D1, D2)` lifts `D1`'s *goal items* into `D2`'s axiom list, which is the right operation only when the two systems agree on an item algebra. Here they do not: the grammar's goal items are `span(0, n, S, phi)` five-tuples, while the prover's rules match `Claim(phi)`. Feeding the former to the latter yields a chart in which no rule can fire. Passing `Claim(phi)` explicitly, weighted by the grammar's goal weight, is the composite that actually typechecks, and it preserves the property that matters: the prover's score is a differentiable function of the grammar's lexicon weights, so gradient descent on an entailment error is gradient descent on the lexicon. + +## Limitations + +The fragment is deliberately small, and two limits are structural rather than incidental. + +First, there is no beta-reduction on the deduction surface, so the grammar cannot use continuation-form determiner denotations and recover a readable logical form; it builds the normal form directly instead. A grammar that genuinely needed higher-order denotations (quantifier raising, for instance, or the scope ambiguities the [quantifier-scope example](quantifier-scope.md) encodes in its categories) would have to encode the reduction in its rules or its categories, not rely on the runtime to normalise. +Second, a rule pattern cannot mention a nullary constant inside a logical form. Rule patterns compile atoms to a tagged `("atom", name)` pair while the lexicon's logical-form evaluator emits a bare `(name,)` tuple, so a pattern like `Claim(Quant(every_t, P, Q))` never matches a chart item. That is why quantificational force is carried by the determiner's category here rather than by a constant in the logical form. The workaround costs one category per determiner, which is why the fragment has `DetEvery` and `DetSome` instead of a single `Det`. ## References -- Richard Montague. 1973. The proper treatment of quantification in ordinary English. In K. J. J. Hintikka, J. M. E. Moravcsik, and P. Suppes, editors, *Approaches to Natural Language*, pages 221–242. Springer, Dordrecht. +- Jon Barwise and Robin Cooper. 1981. [Generalized quantifiers and natural language](https://doi.org/10.1007/BF00350139). *Linguistics and Philosophy*, 4(2):159–219. +- Richard Montague. 1973. [The proper treatment of quantification in ordinary English](https://doi.org/10.1007/978-94-010-2506-5_10). In K. J. J. Hintikka, J. M. E. Moravcsik, and P. Suppes, editors, *Approaches to Natural Language*, pages 221–242. Springer, Dordrecht. diff --git a/docs/examples/seq2seq.md b/docs/examples/seq2seq.md index ca3f63be..51cc720c 100644 --- a/docs/examples/seq2seq.md +++ b/docs/examples/seq2seq.md @@ -14,19 +14,19 @@ object FFHidden, Combined : Real 32 morphism src_embed : Source -> Latent [role=embed] morphism tgt_embed : Target -> Latent [role=embed] -morphism enc_head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal -morphism enc_attn_proj : Latent -> Latent [scale=0.1] ~ Normal -morphism enc_residual_attn : Latent -> Latent [scale=0.01] ~ Normal -morphism enc_ff_up : Latent -> FFHidden ~ Normal -morphism enc_ff_down : FFHidden -> Latent [scale=0.1] ~ Normal -morphism enc_residual_ff : Latent -> Latent [scale=0.01] ~ Normal -morphism dec_head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal -morphism dec_attn_proj : Latent -> Latent [scale=0.1] ~ Normal -morphism dec_residual_attn : Latent -> Latent [scale=0.01] ~ Normal -morphism dec_ff_up : Latent -> FFHidden ~ Normal -morphism dec_ff_down : FFHidden -> Latent [scale=0.1] ~ Normal -morphism dec_residual_ff : Latent -> Latent [scale=0.01] ~ Normal -morphism cross : Combined -> Combined [scale=0.1] ~ Normal +morphism enc_head : Latent -> HeadOut [replicate=4, param_source=mlp] ~ Normal +morphism enc_attn_proj : Latent -> Latent [param_source=mlp] ~ Normal +morphism enc_residual_attn : Latent -> Latent [param_source=mlp] ~ Normal +morphism enc_ff_up : Latent -> FFHidden [param_source=mlp] ~ Normal +morphism enc_ff_down : FFHidden -> Latent [param_source=mlp] ~ Normal +morphism enc_residual_ff : Latent -> Latent [param_source=mlp] ~ Normal +morphism dec_head : Latent -> HeadOut [replicate=4, param_source=mlp] ~ Normal +morphism dec_attn_proj : Latent -> Latent [param_source=mlp] ~ Normal +morphism dec_residual_attn : Latent -> Latent [param_source=mlp] ~ Normal +morphism dec_ff_up : Latent -> FFHidden [param_source=mlp] ~ Normal +morphism dec_ff_down : FFHidden -> Latent [param_source=mlp] ~ Normal +morphism dec_residual_ff : Latent -> Latent [param_source=mlp] ~ Normal +morphism cross : Combined -> Combined [param_source=mlp] ~ Normal morphism lm_head : Combined -> Target ~ Categorical define enc_block = fan(enc_head) >> enc_attn_proj >> enc_residual_attn >> enc_ff_up >> enc_ff_down >> enc_residual_ff diff --git a/docs/examples/source/bidirectional_rnn_lm.qvr b/docs/examples/source/bidirectional_rnn_lm.qvr index c00f84b6..1ae3e2f8 100644 --- a/docs/examples/source/bidirectional_rnn_lm.qvr +++ b/docs/examples/source/bidirectional_rnn_lm.qvr @@ -27,9 +27,9 @@ object Embedded, FwdHidden, BwdHidden : Real 64 object Combined : Real 128 morphism tok_embed : Token -> Embedded [role=embed] -morphism fwd_cell : Embedded * FwdHidden -> FwdHidden [scale=0.1] ~ Normal -morphism bwd_cell : Embedded * BwdHidden -> BwdHidden [scale=0.1] ~ Normal -morphism combine : Combined -> Combined [scale=0.1] ~ Normal +morphism fwd_cell : Embedded * FwdHidden -> FwdHidden [param_source=mlp] ~ Normal +morphism bwd_cell : Embedded * BwdHidden -> BwdHidden [param_source=mlp] ~ Normal +morphism combine : Combined -> Combined [param_source=mlp] ~ Normal morphism lm_head : Combined -> Token ~ Categorical define forward_path = tok_embed >> scan(fwd_cell) diff --git a/docs/examples/source/bnn.qvr b/docs/examples/source/bnn.qvr index 0ac19f11..6e930e42 100644 --- a/docs/examples/source/bnn.qvr +++ b/docs/examples/source/bnn.qvr @@ -1,40 +1,36 @@ -# Bayesian Linear Network for Per-Item Classification Scores +# Bayesian Neural Network for Nonlinear Regression # -# A linear Bayesian network with three weight matrices expressed -# entirely as a morphism composition under the real algebra. The -# per-item input morphism X : Item -> H_in carries the data; -# the final composite scores each Item under H_out output -# channels. +# A Bayesian multi-layer perceptron. The response is Normal about a +# nonlinear function of the input, and that function is a +# two-hidden-layer tanh MLP carried inside the morphism's parameter +# network. The network emits both the mean and the log-scale, so the +# model is heteroscedastic: the predictive spread varies with the +# input. # -# Structural form: +# Generative structure: # -# X : Item -> H_in per-item input -# W_1 : H_in -> H1 matrix-normal weight -# W_2 : H1 -> H2 matrix-normal weight -# W_3 : H2 -> H_out matrix-normal weight -# bnn = X >> W_1 >> W_2 >> W_3 composed linear chain +# mu(x), sigma(x) = MLP(x) two hidden layers, tanh +# y_n ~ Normal(mu(x_n), sigma(x_n)) # -# Each weight tensor is a learnable morphism whose prior is a -# matrix-normal over its (in, out) Kronecker covariance, the -# natural prior for a linear layer. Composition under algebra -# real is strictly linear: there is no pointwise nonlinearity -# between composed weight matrices, so the surface expressed -# here is a Bayesian linear network rather than a deep MLP. - -composition real [level=algebra] - -object Item : FinSet 200 -object H_in : FinSet 4 -object H1 : FinSet 32 -object H2 : FinSet 16 -object H_out : FinSet 2 +# The per-row input x is a free variable, supplied as host data +# through the observations dict; the Resp plate indexes the rows. +# +# What makes the fit Bayesian is a prior over the network's weights. +# The weights are ordinary learnable parameters here, so the prior is +# supplied at the inference layer by lift_from_log_prob rather than +# declared in the source: a Normal prior on every weight turns a +# maximum-likelihood fit into a posterior over networks. +# +# Reference: [MacKay (1992)](https://doi.org/10.1162/neco.1992.4.3.448). -morphism X : Item -> H_in [role=latent] +object Feature : Real 1 +object Target : Real 1 +object Resp : FinSet 200 -morphism W_1 : H_in -> H1 [role=latent] -morphism W_2 : H1 -> H2 [role=latent] -morphism W_3 : H2 -> H_out [role=latent] +morphism net : Feature -> Target [param_source=mlp, hidden_dim=64] ~ Normal -define bnn = X >> W_1 >> W_2 >> W_3 +program bnn : Resp -> Resp + observe y : Resp <- net(x) + return y export bnn diff --git a/docs/examples/source/continuous_hmm.qvr b/docs/examples/source/continuous_hmm.qvr index c82b17c6..03761f0a 100644 --- a/docs/examples/source/continuous_hmm.qvr +++ b/docs/examples/source/continuous_hmm.qvr @@ -20,8 +20,8 @@ object State : Real 16 object Obs : Real 8 -morphism transition : State -> State [scale=0.1] ~ Normal -morphism emission : State -> Obs [scale=0.1] ~ Normal +morphism transition : State -> State ~ Normal +morphism emission : State -> Obs ~ Normal program generative_step : State -> State sample s_new <- transition @@ -29,11 +29,11 @@ program generative_step : State -> State observe o <- emission(s_new) return s_new -morphism inference_cell : Obs * State -> State [scale=0.1] ~ Normal +morphism inference_cell : Obs * State -> State ~ Normal define filter = scan(inference_cell) -morphism decoder : State -> Obs [scale=0.1] ~ Normal +morphism decoder : State -> Obs ~ Normal define filter_and_reconstruct = scan(inference_cell) >> decoder diff --git a/docs/examples/source/deep_markov.qvr b/docs/examples/source/deep_markov.qvr index c416974c..3298ec34 100644 --- a/docs/examples/source/deep_markov.qvr +++ b/docs/examples/source/deep_markov.qvr @@ -1,10 +1,9 @@ # Deep Markov Model # -# A state-space model with nonlinear, -# neural-network-parameterised transition and emission kernels. -# The combinator surface mirrors the linear-Gaussian SSM but -# replaces the linear maps with two-layer MLPs carrying Normal -# weight priors. +# A state-space model with nonlinear, neural-network-parameterised +# transition and emission kernels. The combinator surface mirrors the +# linear-Gaussian SSM and replaces its linear maps with kernels whose +# parameters come from an MLP. # # Generative structure: # @@ -12,6 +11,11 @@ # o_t ~ emit_mlp(s_t) nonlinear emission # s_t ~ infer_cell(o_t, s_{t-1}) variational recognition # +# The nonlinearity is the ``param_source=mlp`` option, which routes a +# kernel's input to its Normal parameters through a tanh MLP. Without +# it a kernel's parameters are linear in its input, and composing two +# of them stays linear: the depth would buy nothing but a rank bound. +# # The inference cell q_phi(o_t, s_{t-1}) -> s_t carries the # variational recognition network; scan threads it across the # observation sequence to recover an amortised posterior over @@ -24,11 +28,11 @@ object Hidden : Real 32 object State : Real 8 object Obs : Real 4 -morphism trans_mlp_1 : Driver * State -> Hidden [scale=0.5] ~ Normal -morphism trans_mlp_2 : Hidden -> State [scale=0.1] ~ Normal -morphism emit_mlp_1 : State -> Hidden [scale=0.5] ~ Normal -morphism emit_mlp_2 : Hidden -> Obs [scale=0.1] ~ Normal -morphism infer_cell : Obs * State -> State [scale=0.1] ~ Normal +morphism trans_mlp_1 : Driver * State -> Hidden [param_source=mlp] ~ Normal +morphism trans_mlp_2 : Hidden -> State [param_source=mlp] ~ Normal +morphism emit_mlp_1 : State -> Hidden [param_source=mlp] ~ Normal +morphism emit_mlp_2 : Hidden -> Obs [param_source=mlp] ~ Normal +morphism infer_cell : Obs * State -> State [param_source=mlp] ~ Normal define transition_cell = trans_mlp_1 >> trans_mlp_2 define emission = emit_mlp_1 >> emit_mlp_2 diff --git a/docs/examples/source/linear_gaussian_ssm.qvr b/docs/examples/source/linear_gaussian_ssm.qvr index 40875581..d598098c 100644 --- a/docs/examples/source/linear_gaussian_ssm.qvr +++ b/docs/examples/source/linear_gaussian_ssm.qvr @@ -25,9 +25,9 @@ object Driver : Real 2 object State : Real 4 object Obs : Real 2 -morphism transition_cell : Driver * State -> State [scale=0.1] ~ Normal -morphism emission : State -> Obs [scale=0.1] ~ Normal -morphism filter_cell : Obs * State -> State [scale=0.1] ~ Normal +morphism transition_cell : Driver * State -> State ~ Normal +morphism emission : State -> Obs ~ Normal +morphism filter_cell : Obs * State -> State ~ Normal define generate = scan(transition_cell) >> emission define filter = scan(filter_cell) diff --git a/docs/examples/source/lstm_lm.qvr b/docs/examples/source/lstm_lm.qvr index 807993dd..1dee5d92 100644 --- a/docs/examples/source/lstm_lm.qvr +++ b/docs/examples/source/lstm_lm.qvr @@ -27,7 +27,7 @@ object Hidden : Real 128 morphism tok_embed : Token -> Embedded [role=embed] morphism gate_i, gate_f, gate_o : Embedded * Hidden -> Hidden ~ LogitNormal -morphism cell_cand : Embedded * Hidden -> Hidden [scale=0.5] ~ Normal +morphism cell_cand : Embedded * Hidden -> Hidden ~ Normal morphism lm_head : Hidden -> Token ~ Categorical program lstm_cell(x_t, h_prev) : Embedded * Hidden -> Hidden diff --git a/docs/examples/source/mixture_model.qvr b/docs/examples/source/mixture_model.qvr index 6ee5bb31..51008d23 100644 --- a/docs/examples/source/mixture_model.qvr +++ b/docs/examples/source/mixture_model.qvr @@ -1,41 +1,38 @@ # Bayesian Gaussian Mixture Model # -# A finite Gaussian mixture model with K components. Per-component -# means and scales are continuous latents on the Component plate; -# the per-row component assignment is a discrete latent integrated -# out by a scoped `marginalize` block under the canonical -# logsumexp aggregation over the K classes. +# A finite Gaussian mixture with K components. Per-component means, +# scales, and mixing weights are latents on the Component plate; each +# observed row is drawn from the resulting mixture. The per-row +# component assignment is integrated out in closed form by the +# MixtureNormal likelihood, so the model carries no discrete latent. # # Generative structure: # -# probs ~ Dirichlet(alpha) Component-simplex mixing weights -# mu_k ~ Normal(0, 5) per-component mean -# sigma_k ~ HalfNormal(1) per-component scale -# cls_n ~ Categorical(probs) per-row class assignment -# r_n ~ Normal(mu[cls_n], sigma[cls_n]) observed row +# probs ~ Dirichlet(alpha) Component-simplex mixing weights +# mu_k ~ Normal(0, 5) per-component mean +# sigma_k ~ HalfNormal(1) per-component scale +# r_n ~ MixtureNormal(probs, mu, sigma) observed row # -# Per-row marginal likelihood (closed form via logsumexp): +# Per-row marginal likelihood (closed form): # # p(r_n) = sum_k probs[k] * Normal(r_n; mu[k], sigma[k]). # -# The `[over=Item]` plate routes each per-row draw through its -# item-group accumulator; `[via=idx]` names the per-observe -# fibration from `Resp` into the `Item` plate. +# MixtureNormal takes three per-component vector parameters (the +# mixing weights, the locations, and the scales), each shared across +# every row of the Resp plate, and scores each row against the +# K-component mixture they define. composition log_prob [level=algebra] object Component : FinSet 3 -object Item : FinSet 8 object Resp : FinSet 100 program gmm(alpha : Real) : Resp -> Resp sample probs <- Dirichlet(alpha) [over=Component] sample mu : Component <- Normal(0.0, 5.0) sample sigma : Component <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) - marginalize cls : Component <- Categorical(probs) [over=Item, reduction=logsumexp] - observe r : Resp <- Normal(mu[cls], sigma[cls]) [via=idx] + observe r : Resp <- MixtureNormal(probs, mu, sigma) return probs diff --git a/docs/examples/source/montague_nli.qvr b/docs/examples/source/montague_nli.qvr index 41cddbec..799e179a 100644 --- a/docs/examples/source/montague_nli.qvr +++ b/docs/examples/source/montague_nli.qvr @@ -1,21 +1,43 @@ -# Montague Grammar + Lambda-Calculus Type Surface +# Montague Grammar + Syllogistic Entailment # # Two-stage natural-language-inference architecture composed # entirely out of QVR's weighted-deduction surface: # -# 1. A Montague-style grammar derives logical forms from token -# spans. Lexical entries are honest lambda terms over a -# user-declared free term algebra; binders are listed in the -# ``binders`` block so the compiler treats their first -# argument as a bound variable and alpha-renames it to a -# fresh canonical symbol per term construction. +# 1. A Montague-style grammar derives a logical form for every +# token span. Common nouns and intransitive verbs denote +# honest lambda terms over a user-declared free term algebra; +# ``Lam`` is listed in the ``binders`` block, so the compiler +# treats its first argument as a bound variable and +# alpha-renames it to a fresh canonical symbol per term +# construction. Structural equality on the chart is therefore +# alpha-equivalence on the surface. # -# 2. An entailment prover closes the LFs under modus ponens and -# a functorial-substitution rule; ``subst(...)`` is the -# capture-avoiding substitution primitive both halves rely -# on. The two deductions chain through ``compose(...)`` at -# the ``program``-body level, so the prover's chart inherits -# the grammar's gradient flow into the lexicon's log-weights. +# 2. An entailment prover closes the resulting claims under +# three syllogistic rules. The prover's items are the +# grammar's own logical forms wrapped in ``Claim``, so the +# two deductions share one term language. +# +# Why the sentence LF is built in normal form +# ------------------------------------------- +# A deduction rule's conclusion is a constructor tree over the +# rule's pattern variables, and the runtime instantiates it by +# structural substitution alone. There is no beta-rule: nothing +# reduces a redex, and the chart never normalises an item. So a +# determiner cannot denote a continuation-form lambda term +# ``Lam(P, Lam(Q, ...))`` that later beta-reduces against its +# arguments; the chart would simply carry the unreduced redex +# forever, and no prover rule could see through it. +# +# The grammar therefore builds the sentence LF in normal form +# directly, in the generalised-quantifier style of +# [Barwise and Cooper (1981)](https://doi.org/10.1007/BF00350139): +# a determiner denotes a relation between two sets, written here +# as the binary constructors ``Every(P, Q)`` and ``Some(P, Q)``. +# The determiner's quantificational force rides on its category +# (``DetEvery`` / ``DetSome``), and the sentence rule that consumes +# that category builds the matching constructor. Restrictor and +# scope stay genuine lambda terms, so the ``binders`` machinery is +# doing real work at the predicate level. # # Reference: [Montague (1973)](https://doi.org/10.1007/978-94-010-2506-5_10). @@ -25,51 +47,85 @@ object Term : FinSet 8 # Grammar deduction # --------------------------------------------------------------------------- -deduction Montague : Term -> Term [semiring=LogProb, start=S, depth=10] - # Category and term-algebra constructors live in ``atoms``. - # ``Var`` is the variable-occurrence constructor; ``App`` is - # application; ``forall_t`` / ``implies_t`` / ``and_t`` are the - # logical connectives used in determiner denotations; predicate - # constants ``dog_p`` / ``bark_p`` / ... are nullary atoms. - atoms NP, S, N, Fwd, Bwd, span, App, Var, dog_p, cat_p, animal_p, bark_p, walk_p, forall_t, exists_t, implies_t, and_t +deduction Montague : Term -> Term [semiring=LogProb, start=S, depth=12] + # Categories, the chart-item constructor ``span``, the LF + # constructors (``App``, ``Var``, ``Every``, ``Some``), and the + # predicate constants all live in ``atoms``. + atoms S, N, VP, Nom, Art, Cop, DetEvery, DetSome, QEvery, QSome, span, App, Var, Every, Some, dog_p, cat_p, animal_p, bark_p, walk_p, every_q, some_q, an_q, is_q # ``Lam`` is the binder constructor: ``Lam(x, body)``. The # compiler alpha-renames ``x`` to a fresh canonical symbol per - # term construction so structural equality on the chart is - # alpha-equivalence on the surface. + # lexicon entry, so the predicate a word denotes is one + # canonical term wherever that word occurs. binders Lam - rule fwd_app : span(I, K, Fwd(X, Y), F), span(K, J, X, A) |- span(I, J, Y, App(F, A)) #[learnable] - rule bwd_app : span(I, K, X, A), span(K, J, Bwd(X, Y), F) |- span(I, J, Y, App(F, A)) #[learnable] + # Determiner + common noun. The quantified-NP span carries the + # restrictor predicate as its LF; the determiner's category + # records which quantifier is coming. + rule every_np : span(I, K, DetEvery, D), span(K, J, N, P) |- span(I, J, QEvery, P) #[learnable] + rule some_np : span(I, K, DetSome, D), span(K, J, N, P) |- span(I, J, QSome, P) #[learnable] + # Quantified NP + verb phrase. This is where the + # generalised-quantifier relation is built, already in normal + # form: no redex, nothing left to reduce. + rule every_s : span(I, K, QEvery, P), span(K, J, VP, Q) |- span(I, J, S, Every(P, Q)) #[learnable] + rule some_s : span(I, K, QSome, P), span(K, J, VP, Q) |- span(I, J, S, Some(P, Q)) #[learnable] + # Predicative "is an N": the article and the copula are + # semantically vacuous, so both rules pass the noun's predicate + # through unchanged. + rule art_n : span(I, K, Art, D), span(K, J, N, P) |- span(I, J, Nom, P) #[learnable] + rule cop_nom : span(I, K, Cop, D), span(K, J, Nom, P) |- span(I, J, VP, P) #[learnable] lexicon # Common nouns are unary predicates: ``λx. dog(x)``. - "dog" : N = Lam(x, App(dog_p, Var(x))) #[learnable] - "cat" : N = Lam(x, App(cat_p, Var(x))) #[learnable] - "animal" : N = Lam(x, App(animal_p, Var(x))) #[learnable] + "dog" : N = Lam(x, App(dog_p, Var(x))) #[learnable] + "cat" : N = Lam(x, App(cat_p, Var(x))) #[learnable] + "animal" : N = Lam(x, App(animal_p, Var(x))) #[learnable] # Intransitive verbs are unary predicates over individuals. - "barks" : Bwd(NP, S) = Lam(x, App(bark_p, Var(x))) #[learnable] - "walks" : Bwd(NP, S) = Lam(x, App(walk_p, Var(x))) #[learnable] - # Determiners are generalised quantifiers: - # "every" : (S/(S\NP))/N - # = λP. λQ. ∀(λx. P(x) → Q(x)) - # Bound variables ``P``, ``Q``, ``x`` are declared by the - # ``binders Lam`` block + the lexicon-LF compiler's bound- - # variable pre-pass; they need not appear in ``atoms``. - "every" : Fwd(N, Fwd(Bwd(NP, S), S)) = Lam(P, Lam(Q, App(forall_t, Lam(x, App(App(implies_t, App(Var(P), Var(x))), App(Var(Q), Var(x))))))) #[learnable] - "some" : Fwd(N, Fwd(Bwd(NP, S), S)) = Lam(P, Lam(Q, App(exists_t, Lam(x, App(App(and_t, App(Var(P), Var(x))), App(Var(Q), Var(x))))))) #[learnable] + "barks" : VP = Lam(x, App(bark_p, Var(x))) #[learnable] + "walks" : VP = Lam(x, App(walk_p, Var(x))) #[learnable] + # The determiners, the article, and the copula contribute no + # LF of their own: their category is what drives the rules, + # exactly as in the CCG example's lexicon. Each entry names + # itself with a constant so its log-weight is still learnable. + "every" : DetEvery = every_q #[learnable] + "some" : DetSome = some_q #[learnable] + "an" : Art = an_q #[learnable] + "is" : Cop = is_q #[learnable] # --------------------------------------------------------------------------- # Prover deduction # --------------------------------------------------------------------------- +# +# The prover's items are ``Claim(phi)`` for ``phi`` a sentence LF +# the grammar built, plus ``Nonempty(P)`` for an explicit +# existence premise. All three rules are classically valid on the +# generalised-quantifier reading +# +# Every(P, Q) = forall x. P(x) -> Q(x) +# Some(P, Q) = exists x. P(x) & Q(x) +# Nonempty(P) = exists x. P(x) +# +# ``ex_import`` is the one that needs care. ``Every(P, Q)`` alone +# does *not* entail ``Some(P, Q)``: on the standard first-order +# reading a universal over an empty restrictor is vacuously true, +# while the existential is false. The inference is valid only +# under existential import, so the rule takes ``Nonempty(P)`` as a +# second premise rather than smuggling the assumption in. A caller +# that wants the inference must supply that premise. -deduction Prover : Term -> Term [semiring=LogProb, depth=8] - atoms Claim, Implies, App, Var +deduction Prover : Term -> Term [semiring=LogProb, depth=12] + atoms Claim, Every, Some, Nonempty, App, Var binders Lam - rule modus_ponens : Claim(P), Claim(Implies(P, Q)) |- Claim(Q) #[learnable] - rule app_subst : Claim(App(F, X)), Claim(Implies(X, Y)) |- Claim(App(F, Y)) #[learnable] + # Barbara: every P is Q, every Q is R, therefore every P is R. + # Valid with no existence assumption. + rule barbara : Claim(Every(P, Q)), Claim(Every(Q, R)) |- Claim(Every(P, R)) #[learnable] + # Darii: some P is Q, every Q is R, therefore some P is R. + # Valid; the existential premise carries the witness. + rule darii : Claim(Some(P, Q)), Claim(Every(Q, R)) |- Claim(Some(P, R)) #[learnable] + # Subalternation, licensed only by an explicit non-emptiness + # premise on the restrictor. + rule ex_import : Claim(Every(P, Q)), Claim(Nonempty(P)) |- Claim(Some(P, Q)) #[learnable] # --------------------------------------------------------------------------- -# Fit program: regression-style SVI / NUTS over the lexicon and -# rule log-weights. The chart's goal weight is the inside -# log-marginal of the sentence under the current parameters. +# Fit program: the chart's goal weight is the inside log-marginal +# of the sentence at the start symbol under the current parameters. # --------------------------------------------------------------------------- program fit_grammar : Term -> Term @@ -78,18 +134,3 @@ program fit_grammar : Term -> Term return log_Z export fit_grammar - -# --------------------------------------------------------------------------- -# Demonstration of ``subst`` and ``compose`` as program-body -# builtins. ``subst(term, var, value)`` is capture-avoiding because -# every binder has already been alpha-renamed at compile time; -# ``compose(D1, D2)`` chains D1's goal items into D2's axiom set -# and propagates submodule parameters so a composed fit walks both -# factors' weights together. -# --------------------------------------------------------------------------- - -program fit_pipeline : Term -> Term - let pipeline = compose(Montague, Prover) - let chart = parse(pipeline, sentence) - score log_Z = chart.goal_weight() - return log_Z diff --git a/docs/examples/source/seq2seq.qvr b/docs/examples/source/seq2seq.qvr index cf31fd55..226c8cf0 100644 --- a/docs/examples/source/seq2seq.qvr +++ b/docs/examples/source/seq2seq.qvr @@ -28,19 +28,19 @@ object FFHidden, Combined : Real 32 morphism src_embed : Source -> Latent [role=embed] morphism tgt_embed : Target -> Latent [role=embed] -morphism enc_head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal -morphism enc_attn_proj : Latent -> Latent [scale=0.1] ~ Normal -morphism enc_residual_attn : Latent -> Latent [scale=0.01] ~ Normal -morphism enc_ff_up : Latent -> FFHidden ~ Normal -morphism enc_ff_down : FFHidden -> Latent [scale=0.1] ~ Normal -morphism enc_residual_ff : Latent -> Latent [scale=0.01] ~ Normal -morphism dec_head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal -morphism dec_attn_proj : Latent -> Latent [scale=0.1] ~ Normal -morphism dec_residual_attn : Latent -> Latent [scale=0.01] ~ Normal -morphism dec_ff_up : Latent -> FFHidden ~ Normal -morphism dec_ff_down : FFHidden -> Latent [scale=0.1] ~ Normal -morphism dec_residual_ff : Latent -> Latent [scale=0.01] ~ Normal -morphism cross : Combined -> Combined [scale=0.1] ~ Normal +morphism enc_head : Latent -> HeadOut [replicate=4, param_source=mlp] ~ Normal +morphism enc_attn_proj : Latent -> Latent [param_source=mlp] ~ Normal +morphism enc_residual_attn : Latent -> Latent [param_source=mlp] ~ Normal +morphism enc_ff_up : Latent -> FFHidden [param_source=mlp] ~ Normal +morphism enc_ff_down : FFHidden -> Latent [param_source=mlp] ~ Normal +morphism enc_residual_ff : Latent -> Latent [param_source=mlp] ~ Normal +morphism dec_head : Latent -> HeadOut [replicate=4, param_source=mlp] ~ Normal +morphism dec_attn_proj : Latent -> Latent [param_source=mlp] ~ Normal +morphism dec_residual_attn : Latent -> Latent [param_source=mlp] ~ Normal +morphism dec_ff_up : Latent -> FFHidden [param_source=mlp] ~ Normal +morphism dec_ff_down : FFHidden -> Latent [param_source=mlp] ~ Normal +morphism dec_residual_ff : Latent -> Latent [param_source=mlp] ~ Normal +morphism cross : Combined -> Combined [param_source=mlp] ~ Normal morphism lm_head : Combined -> Target ~ Categorical define enc_block = fan(enc_head) >> enc_attn_proj >> enc_residual_attn >> enc_ff_up >> enc_ff_down >> enc_residual_ff diff --git a/docs/examples/source/transformer_lm.qvr b/docs/examples/source/transformer_lm.qvr index 14f944d6..0f44c7a1 100644 --- a/docs/examples/source/transformer_lm.qvr +++ b/docs/examples/source/transformer_lm.qvr @@ -26,11 +26,11 @@ object HeadOut : Real 4 object FFHidden : Real 32 morphism tok_embed : Token -> Latent [role=embed] -morphism head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal -morphism attn_proj : Latent -> Latent [scale=0.1] ~ Normal -morphism ff_up : Latent -> FFHidden ~ Normal -morphism ff_down : FFHidden -> Latent [scale=0.1] ~ Normal -morphism residual_attn, residual_ff : Latent -> Latent [scale=0.01] ~ Normal +morphism head : Latent -> HeadOut [replicate=4, param_source=mlp] ~ Normal +morphism attn_proj : Latent -> Latent [param_source=mlp] ~ Normal +morphism ff_up : Latent -> FFHidden [param_source=mlp] ~ Normal +morphism ff_down : FFHidden -> Latent [param_source=mlp] ~ Normal +morphism residual_attn, residual_ff : Latent -> Latent [param_source=mlp] ~ Normal morphism lm_head : Latent -> Token ~ Categorical define layer = fan(head) >> attn_proj >> residual_attn >> ff_up >> ff_down >> residual_ff diff --git a/docs/examples/source/vae.qvr b/docs/examples/source/vae.qvr index 8b9b487e..18869c34 100644 --- a/docs/examples/source/vae.qvr +++ b/docs/examples/source/vae.qvr @@ -28,16 +28,16 @@ object ObsSpace : Real 8 object UnitSpace : Real 1 morphism pixel_embed : Pixel -> EncoderHidden [role=embed] -morphism enc_deep : EncoderHidden -> EncoderHidden ~ Normal -morphism enc_to_latent : EncoderHidden -> Latent [scale=0.5] ~ Normal +morphism enc_deep : EncoderHidden -> EncoderHidden [param_source=mlp] ~ Normal +morphism enc_to_latent : EncoderHidden -> Latent ~ Normal define encoder = pixel_embed >> stack(enc_deep, 1) >> enc_to_latent morphism prior : UnitSpace -> Latent ~ Normal morphism dec_1 : Latent -> DecoderHidden ~ Normal -morphism dec_deep : DecoderHidden -> DecoderHidden ~ Normal -morphism dec_to_obs : DecoderHidden -> ObsSpace [scale=0.1] ~ Normal +morphism dec_deep : DecoderHidden -> DecoderHidden [param_source=mlp] ~ Normal +morphism dec_to_obs : DecoderHidden -> ObsSpace ~ Normal define decoder = dec_1 >> stack(dec_deep, 1) >> dec_to_obs diff --git a/docs/examples/source/vanilla_rnn_lm.qvr b/docs/examples/source/vanilla_rnn_lm.qvr index 298f6348..2eb1b199 100644 --- a/docs/examples/source/vanilla_rnn_lm.qvr +++ b/docs/examples/source/vanilla_rnn_lm.qvr @@ -20,7 +20,7 @@ object Embedded : Real 64 object Hidden : Real 128 morphism tok_embed : Token -> Embedded [role=embed] -morphism cell : Embedded * Hidden -> Hidden [scale=0.1] ~ Normal +morphism cell : Embedded * Hidden -> Hidden [param_source=mlp] ~ Normal morphism lm_head : Hidden -> Token ~ Categorical define backbone = tok_embed >> scan(cell) diff --git a/docs/examples/transformer-lm.md b/docs/examples/transformer-lm.md index 303e3774..4a5a5351 100644 --- a/docs/examples/transformer-lm.md +++ b/docs/examples/transformer-lm.md @@ -13,11 +13,11 @@ object HeadOut : Real 4 object FFHidden : Real 32 morphism tok_embed : Token -> Latent [role=embed] -morphism head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal -morphism attn_proj : Latent -> Latent [scale=0.1] ~ Normal -morphism ff_up : Latent -> FFHidden ~ Normal -morphism ff_down : FFHidden -> Latent [scale=0.1] ~ Normal -morphism residual_attn, residual_ff : Latent -> Latent [scale=0.01] ~ Normal +morphism head : Latent -> HeadOut [replicate=4, param_source=mlp] ~ Normal +morphism attn_proj : Latent -> Latent [param_source=mlp] ~ Normal +morphism ff_up : Latent -> FFHidden [param_source=mlp] ~ Normal +morphism ff_down : FFHidden -> Latent [param_source=mlp] ~ Normal +morphism residual_attn, residual_ff : Latent -> Latent [param_source=mlp] ~ Normal morphism lm_head : Latent -> Token ~ Categorical define layer = fan(head) >> attn_proj >> residual_attn >> ff_up >> ff_down >> residual_ff @@ -36,7 +36,7 @@ export transformer_lm ### Multi-head attention -`morphism head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal` declares four independent attention heads via the [replicate](../guides/dsl-declarations.md#replicated-declarations) attribute on a single morphism. Each head is a Bayesian Kleisli morphism `Latent -> HeadOut`; `HeadOut` is four-dimensional, so the four heads together cover the sixteen-dimensional `Latent`. [`fan(head)`](../guides/dsl-declarations.md#fan-out-diagonal-morphism) runs the four heads in parallel on the same input and concatenates the outputs, the standard multi-head wiring. +`morphism head : Latent -> HeadOut [replicate=4, param_source=mlp] ~ Normal` declares four independent attention heads via the [replicate](../guides/dsl-declarations.md#replicated-declarations) attribute on a single morphism. Each head is a Bayesian Kleisli morphism `Latent -> HeadOut`; `HeadOut` is four-dimensional, so the four heads together cover the sixteen-dimensional `Latent`. [`fan(head)`](../guides/dsl-declarations.md#fan-out-diagonal-morphism) runs the four heads in parallel on the same input and concatenates the outputs, the standard multi-head wiring. ### Layer block diff --git a/docs/examples/vae.md b/docs/examples/vae.md index c4a5eea1..393e4a96 100644 --- a/docs/examples/vae.md +++ b/docs/examples/vae.md @@ -14,16 +14,16 @@ object ObsSpace : Real 8 object UnitSpace : Real 1 morphism pixel_embed : Pixel -> EncoderHidden [role=embed] -morphism enc_deep : EncoderHidden -> EncoderHidden ~ Normal -morphism enc_to_latent : EncoderHidden -> Latent [scale=0.5] ~ Normal +morphism enc_deep : EncoderHidden -> EncoderHidden [param_source=mlp] ~ Normal +morphism enc_to_latent : EncoderHidden -> Latent ~ Normal define encoder = pixel_embed >> stack(enc_deep, 1) >> enc_to_latent morphism prior : UnitSpace -> Latent ~ Normal morphism dec_1 : Latent -> DecoderHidden ~ Normal -morphism dec_deep : DecoderHidden -> DecoderHidden ~ Normal -morphism dec_to_obs : DecoderHidden -> ObsSpace [scale=0.1] ~ Normal +morphism dec_deep : DecoderHidden -> DecoderHidden [param_source=mlp] ~ Normal +morphism dec_to_obs : DecoderHidden -> ObsSpace ~ Normal define decoder = dec_1 >> stack(dec_deep, 1) >> dec_to_obs diff --git a/docs/examples/vanilla-rnn-lm.md b/docs/examples/vanilla-rnn-lm.md index f6e87bcc..40271331 100644 --- a/docs/examples/vanilla-rnn-lm.md +++ b/docs/examples/vanilla-rnn-lm.md @@ -12,7 +12,7 @@ object Embedded : Real 64 object Hidden : Real 128 morphism tok_embed : Token -> Embedded [role=embed] -morphism cell : Embedded * Hidden -> Hidden [scale=0.1] ~ Normal +morphism cell : Embedded * Hidden -> Hidden [param_source=mlp] ~ Normal morphism lm_head : Hidden -> Token ~ Categorical define backbone = tok_embed >> scan(cell) diff --git a/docs/getting-started/highlighting.md b/docs/getting-started/highlighting.md index 85ac63e4..ee65d9ee 100644 --- a/docs/getting-started/highlighting.md +++ b/docs/getting-started/highlighting.md @@ -203,7 +203,7 @@ object SD : Real 32 object SK : Real 64 morphism W : SD -> SK [role=latent, over=[dom, cod]] ~ MatrixNormal(0.0, 1.0) -morphism softmax_link : SK -> SK [scale=0.1] ~ Normal +morphism softmax_link : SK -> SK ~ Normal program regression : SD -> SK [effects = [Sample, Score]] let z = W >> softmax_link observe y : K <- Normal(z, 0.5) diff --git a/docs/guides/analysis-data-and-formulas.md b/docs/guides/analysis-data-and-formulas.md index 06f1e4f1..c60484af 100644 --- a/docs/guides/analysis-data-and-formulas.md +++ b/docs/guides/analysis-data-and-formulas.md @@ -171,13 +171,34 @@ Custom families are pluggable: subclass [`Family`](../api/formulas/family.md#quivers.formulas.family.Family) and register your own observe kernel and link. +### Coefficient priors are autoscaled + +A column enters the linear predictor as `beta * column`, so a prior on +the coefficient alone is really a statement about the coefficient's +contribution, and the same nominal prior means something different for +every column. The default fixed-effect prior is therefore autoscaled: +its scale is divided by the column's root-mean-square, which states it +in contribution space so that `Normal(0.0, 5.0)` means the same thing +on a raw predictor and on an orthonormal `poly` column. The +coefficients themselves stay on their own column's scale, so nothing +needs transforming back. + +This matters most for a basis whose columns are not O(1). `poly(x, k)` +returns columns of norm one, whose entries run about $1/\sqrt{N}$; an +unscaled `Normal(0.0, 5.0)` would assert that the contribution is near +zero, and the fit would agree with the prior rather than the data, +putting the noise scale at the marginal spread of the response and +leaving the coefficients where they started. + ### Prior overrides Prior overrides are keyed by the latent's name in the emitted QVR program (which `formula_to_qvr` lets you inspect upfront). The prior template is a brms-style `Family(arg, arg, ...)` call; numeric args become floats, identifier args stay as references to -other latents in the program. The full call shape lives in +other latents in the program. An explicit prior is your statement +about that coefficient and is emitted exactly as written, without the +autoscaling above. The full call shape lives in [Fitting and Diagnostics](analysis-fitting-and-diagnostics.md#prior-overrides). ## See also diff --git a/docs/guides/dsl-declarations.md b/docs/guides/dsl-declarations.md index e049ccf4..0c2a98f1 100644 --- a/docs/guides/dsl-declarations.md +++ b/docs/guides/dsl-declarations.md @@ -235,7 +235,7 @@ morphism cat : X -> Y * Z # Parametric kernel: input-conditional Normal on R^3. morphism f : X -> R3 ~ Normal # Family options control the parameter network. -morphism g : R3 -> R3 ~ Normal [scale=0.5] +morphism g : R3 -> R3 ~ Normal morphism k : X -> S3 ~ Dirichlet # 30+ families are registered; see the families guide. # `Flow` is a special-case constructor (not in the family registry) @@ -401,7 +401,7 @@ referenced by [`fan`](#fan-out-diagonal-morphism): ```qvr # creates head_0, head_1, head_2, head_3 with independent parameters -morphism head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal +morphism head : Latent -> HeadOut [replicate=4] ~ Normal # works on every morphism whose role permits replication: morphism T : State -> Obs [replicate=3] # lookup-table kernel @@ -427,12 +427,12 @@ declaration: define parallel = fan(f, g, h) # group expansion: fan(head) expands to fan(head_0, head_1, head_2, head_3) -morphism head : Latent -> HeadOut [replicate=4, scale=0.1] ~ Normal +morphism head : Latent -> HeadOut [replicate=4] ~ Normal define multi_head = fan(head) # commonly followed by a projection to recombine -morphism proj : Combined -> Latent [scale=0.1] ~ Normal +morphism proj : Combined -> Latent ~ Normal define attention = fan(head) >> proj ``` @@ -511,7 +511,7 @@ Thread hidden state across a sequence using a recurrent cell: ```qvr # Basic syntax: cell has product domain A * H -> H -morphism cell : Embedded * Hidden -> Hidden ~ Normal [scale=0.1] +morphism cell : Embedded * Hidden -> Hidden ~ Normal define rnn = tok_embed >> scan(cell) >> output_proj # With learned initial state (default is zeros) @@ -549,8 +549,8 @@ object Output : Real 64 morphism tok_embed : Token -> Embedded [role=embed] -morphism cell : Embedded * Hidden -> Hidden ~ Normal [scale=0.1] -morphism output_proj : Hidden -> Output ~ Normal [scale=0.1] +morphism cell : Embedded * Hidden -> Hidden ~ Normal +morphism output_proj : Hidden -> Output ~ Normal define rnn = tok_embed >> scan(cell) >> output_proj export rnn diff --git a/docs/guides/dsl-programs-and-lets.md b/docs/guides/dsl-programs-and-lets.md index f6a66c5c..dd85341a 100644 --- a/docs/guides/dsl-programs-and-lets.md +++ b/docs/guides/dsl-programs-and-lets.md @@ -517,7 +517,7 @@ object Latent : Real 3 object Obs : Real 5 morphism prior : Cond -> Latent ~ Normal -morphism likelihood : Latent -> Obs ~ Normal [scale=0.1] +morphism likelihood : Latent -> Obs ~ Normal define posterior = prior >> likelihood export posterior diff --git a/docs/tutorials/qvr/05-time-series.md b/docs/tutorials/qvr/05-time-series.md index ff940992..fdf35512 100644 --- a/docs/tutorials/qvr/05-time-series.md +++ b/docs/tutorials/qvr/05-time-series.md @@ -39,8 +39,8 @@ object Output : Real 64 morphism tok_embed : Token -> Embedded [role=embed] -morphism cell : Embedded * Hidden -> Hidden [scale=0.1] ~ Normal -morphism output_proj : Hidden -> Output [scale=0.1] ~ Normal +morphism cell : Embedded * Hidden -> Hidden ~ Normal +morphism output_proj : Hidden -> Output ~ Normal define rnn = tok_embed >> scan(cell) >> output_proj export rnn ``` @@ -101,9 +101,9 @@ object Driver : Real 2 object State : Real 4 object Obs : Real 2 -morphism transition_cell : Driver * State -> State [scale=0.1] ~ Normal -morphism emission : State -> Obs [scale=0.1] ~ Normal -morphism filter_cell : Obs * State -> State [scale=0.1] ~ Normal +morphism transition_cell : Driver * State -> State ~ Normal +morphism emission : State -> Obs ~ Normal +morphism filter_cell : Obs * State -> State ~ Normal define generate = scan(transition_cell) >> emission define filter = scan(filter_cell) diff --git a/mkdocs.yml b/mkdocs.yml index 4806d77d..a4956eb8 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -253,6 +253,7 @@ nav: - Overview: api/continuous/spaces.md - Morphisms: api/continuous/morphisms.md - Families: api/continuous/families.md + - Parameter Sources: api/continuous/param_source.md - Inline Builders: api/continuous/inline.md - Programs: api/continuous/programs.md - Scan: api/continuous/scan.md diff --git a/pyproject.toml b/pyproject.toml index 4d46fede..89da7366 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "quivers" -version = "0.15.0" +version = "0.16.0" description = "A functional probabilistic programming language that compiles to PyTorch." readme = "README.md" license = {text = "MIT"} diff --git a/regression.qvr b/regression.qvr index f3b6c956..1c30fbe6 100644 --- a/regression.qvr +++ b/regression.qvr @@ -2,7 +2,7 @@ object Resp : FinSet 200 program model : Resp -> Resp sample intercept <- Normal(0.0, 5.0) - sample beta_x <- Normal(0.0, 5.0) + sample beta_x <- Normal(0.0, 4.883378911472473) let beta_x_per_row = (beta_x * x) sample sigma <- HalfCauchy(2.0) let eta = (intercept + beta_x_per_row) diff --git a/src/quivers/__init__.py b/src/quivers/__init__.py index b1ba961e..708c5bb4 100644 --- a/src/quivers/__init__.py +++ b/src/quivers/__init__.py @@ -20,7 +20,12 @@ output = program() # tensor of shape (3, 2) """ -__version__ = "0.14.1" +from importlib.metadata import version as _version + +#: Read from the installed distribution's metadata, which the build +#: takes from ``pyproject.toml``. A literal here would be a second +#: place to state the version and a second place to forget it. +__version__ = _version("quivers") from quivers.core.objects import ( SetObject, diff --git a/src/quivers/continuous/families.py b/src/quivers/continuous/families.py index ac481bf8..8ee94fed 100644 --- a/src/quivers/continuous/families.py +++ b/src/quivers/continuous/families.py @@ -27,7 +27,7 @@ from __future__ import annotations import math -from collections.abc import Callable +from collections.abc import Callable, Sequence import torch import torch.nn.functional as F @@ -55,8 +55,8 @@ from quivers.continuous.morphisms import ( AnySpace, ContinuousMorphism, - _make_source, ) +from quivers.continuous.param_source import ParamSource, _make_source from quivers.core._util import EPS @@ -125,7 +125,7 @@ def __init__( codomain: ContinuousSpace, dist_class: type, param_specs: list[tuple[str, Callable]], - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, discrete: bool = False, param_source=None, param_source_option: str | None = None, @@ -139,21 +139,12 @@ def __init__( # total raw parameters: one scalar per spec per codomain dim total_raw = len(param_specs) * d - if param_source is None and param_source_option is not None: - from quivers.continuous.param_source import ( - param_source_from_option, - ) - - param_source = param_source_from_option( - domain, - total_raw, - param_source_option, - ) self.param_source = _make_source( domain, total_raw, hidden_dim, param_source=param_source, + param_source_option=param_source_option, ) def _get_dist(self, x: torch.Tensor) -> D.Distribution: @@ -263,8 +254,8 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, - param_source=None, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, param_source_option: str | None = None, ) -> None: super().__init__( @@ -341,28 +332,19 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, param_source=None, param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim # param_dim = d (mu) + d (log_sigma) - if param_source is None and param_source_option is not None: - from quivers.continuous.param_source import ( - param_source_from_option, - ) - - param_source = param_source_from_option( - domain, - 2 * d, - param_source_option, - ) self.param_source = _make_source( domain, 2 * d, hidden_dim, param_source=param_source, + param_source_option=param_source_option, ) self._d = d @@ -433,11 +415,19 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim - self.param_source = _make_source(domain, 2 * d, hidden_dim) + self.param_source = _make_source( + domain, + 2 * d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d @property @@ -511,11 +501,19 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim - self.param_source = _make_source(domain, 2 * d, hidden_dim) + self.param_source = _make_source( + domain, + 2 * d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d @property @@ -571,14 +569,22 @@ def __init__( self, domain: AnySpace, codomain: Euclidean, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: if codomain.low is None or codomain.high is None: raise ValueError("ConditionalTruncatedNormal requires a bounded codomain") super().__init__(domain, codomain) d = codomain.dim - self.param_source = _make_source(domain, 2 * d, hidden_dim) + self.param_source = _make_source( + domain, + 2 * d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d self._low = codomain.low self._high = codomain.high @@ -658,11 +664,19 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim - self.param_source = _make_source(domain, d, hidden_dim) + self.param_source = _make_source( + domain, + d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d @property @@ -875,12 +889,20 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim # param_dim = d (loc) + d (raw_width) - self.param_source = _make_source(domain, 2 * d, hidden_dim) + self.param_source = _make_source( + domain, + 2 * d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d # The bounds are data-dependent, so we cannot pin a single # interval at construction time; advertise the codomain's @@ -945,12 +967,20 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim n_tril = d * (d + 1) // 2 - self.param_source = _make_source(domain, d + n_tril, hidden_dim) + self.param_source = _make_source( + domain, + d + n_tril, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d self._n_tril = n_tril @@ -1017,7 +1047,9 @@ def __init__( domain: AnySpace, codomain: ContinuousSpace, rank: int = 2, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim @@ -1026,7 +1058,13 @@ def __init__( # loc (d) + factor (d * rank) + diag (d) total = d + d * rank + d - self.param_source = _make_source(domain, total, hidden_dim) + self.param_source = _make_source( + domain, + total, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) def _get_dist(self, x: torch.Tensor) -> D.LowRankMultivariateNormal: raw = self.param_source(x) @@ -1084,11 +1122,19 @@ def __init__( domain: AnySpace, codomain: ContinuousSpace, temperature: float = 0.5, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim - self.param_source = _make_source(domain, d, hidden_dim) + self.param_source = _make_source( + domain, + d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d self._temperature = temperature @@ -1137,11 +1183,19 @@ def __init__( domain: AnySpace, codomain: ContinuousSpace, temperature: float = 0.5, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim - self.param_source = _make_source(domain, d, hidden_dim) + self.param_source = _make_source( + domain, + d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d self._temperature = temperature @@ -1195,13 +1249,21 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim n_tril = d * (d + 1) // 2 # df (1) + lower-triangular scale (n_tril) - self.param_source = _make_source(domain, 1 + n_tril, hidden_dim) + self.param_source = _make_source( + domain, + 1 + n_tril, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d self._n_tril = n_tril @@ -1301,7 +1363,9 @@ def __init__( codomain: ContinuousSpace, rows: int, cols: int, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) self._rows = int(rows) @@ -1313,7 +1377,11 @@ def __init__( self._n_row_tril = n_row_tril self._n_col_tril = n_col_tril self.param_source = _make_source( - domain, n_loc + n_row_tril + n_col_tril, hidden_dim + domain, + n_loc + n_row_tril + n_col_tril, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, ) def _build_tril(self, raw: torch.Tensor, d: int, n_tril: int) -> torch.Tensor: @@ -1390,14 +1458,22 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim n_tril = d * (d + 1) // 2 self._d = d self._n_tril = n_tril - self.param_source = _make_source(domain, 1 + n_tril, hidden_dim) + self.param_source = _make_source( + domain, + 1 + n_tril, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -1495,7 +1571,9 @@ def __init__( self, domain: AnySpace, codomain: AnySpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: from quivers.core.objects import SetObject @@ -1507,7 +1585,13 @@ def __init__( super().__init__(domain, codomain) # one logit per input - self.param_source = _make_source(domain, 1, hidden_dim) + self.param_source = _make_source( + domain, + 1, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) def _get_probs(self, x: torch.Tensor) -> torch.Tensor: """Compute Bernoulli probabilities from input. @@ -1589,7 +1673,9 @@ def __init__( self, domain: AnySpace, codomain: AnySpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: from quivers.core.objects import SetObject @@ -1600,7 +1686,13 @@ def __init__( super().__init__(domain, codomain) self._k = codomain.size - self.param_source = _make_source(domain, self._k, hidden_dim) + self.param_source = _make_source( + domain, + self._k, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self): @@ -1692,7 +1784,9 @@ def __init__( domain: AnySpace, codomain: ContinuousSpace, total_count: int = 1, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: if total_count < 1: raise ValueError( @@ -1702,7 +1796,13 @@ def __init__( d = codomain.dim self._d = d self._total_count = int(total_count) - self.param_source = _make_source(domain, d, hidden_dim) + self.param_source = _make_source( + domain, + d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -1749,14 +1849,22 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim # We use a Normal in (d-1)-dim space and the # StickBreakingTransform to land on the d-simplex. # torch.distributions.LogisticNormal handles this. - self.param_source = _make_source(domain, 2 * (d - 1), hidden_dim) + self.param_source = _make_source( + domain, + 2 * (d - 1), + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d @property @@ -1801,7 +1909,9 @@ def __init__( self, domain: AnySpace, codomain: AnySpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: from quivers.core.objects import SetObject @@ -1817,7 +1927,13 @@ def __init__( ) super().__init__(domain, codomain) self._k = codomain.size - self.param_source = _make_source(domain, 1 + (self._k - 1), hidden_dim) + self.param_source = _make_source( + domain, + 1 + (self._k - 1), + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -1856,12 +1972,20 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim self._d = d - self.param_source = _make_source(domain, 2 * d, hidden_dim) + self.param_source = _make_source( + domain, + 2 * d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -1896,12 +2020,20 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim self._d = d - self.param_source = _make_source(domain, 2 * d, hidden_dim) + self.param_source = _make_source( + domain, + 2 * d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -1943,7 +2075,9 @@ def __init__( domain: AnySpace, codomain: ContinuousSpace, num_components: int = 2, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: if num_components < 2: raise ValueError( @@ -1952,7 +2086,13 @@ def __init__( ) super().__init__(domain, codomain) self._k = int(num_components) - self.param_source = _make_source(domain, 3 * self._k, hidden_dim) + self.param_source = _make_source( + domain, + 3 * self._k, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -2002,12 +2142,20 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim self._d = d - self.param_source = _make_source(domain, d, hidden_dim) + self.param_source = _make_source( + domain, + d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -2073,7 +2221,9 @@ def __init__( codomain: ContinuousSpace, component_class: type, num_components: int = 4, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: if num_components < 2: raise ValueError( @@ -2084,7 +2234,13 @@ def __init__( self._components = torch.nn.ModuleList( [component_class(domain, codomain, hidden_dim) for _ in range(self._K)] ) - self.mixture_logits = _make_source(domain, self._K, hidden_dim) + self.mixture_logits = _make_source( + domain, + self._K, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self): # type: ignore[override] @@ -2255,11 +2411,19 @@ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) self._matrix_dim = codomain.dim - self.param_source = _make_source(domain, 1, hidden_dim) + self.param_source = _make_source( + domain, + 1, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) @property def support(self) -> _constraints.Constraint: @@ -2668,20 +2832,31 @@ class ConditionalGeneralizedPareto(ContinuousMorphism): Source space. codomain : ContinuousSpace Target space. - hidden_dim : int - Hidden layer width for neural parameter source. + hidden_dim : int or sequence of int + Hidden widths, read only by a source that has hidden layers. + param_source, param_source_option + Select the parameter source directly, or by the DSL's + ``[param_source=...]`` text. """ def __init__( self, domain: AnySpace, codomain: ContinuousSpace, - hidden_dim: int = 64, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, ) -> None: super().__init__(domain, codomain) d = codomain.dim # loc + scale + concentration - self.param_source = _make_source(domain, 3 * d, hidden_dim) + self.param_source = _make_source( + domain, + 3 * d, + hidden_dim, + param_source=param_source, + param_source_option=param_source_option, + ) self._d = d def _get_dist(self, x: torch.Tensor) -> D.GeneralizedPareto: diff --git a/src/quivers/continuous/inline.py b/src/quivers/continuous/inline.py index ebbe212a..9dc223ae 100644 --- a/src/quivers/continuous/inline.py +++ b/src/quivers/continuous/inline.py @@ -203,11 +203,15 @@ def __init__( self._dist_builder = dist_builder self._discrete = discrete self._support = support if support is not None else _constraints.real - # Per-position event rank. Position 0 = per-row scalar; > 0 = + # Per-position event rank. Position 0 = per-row scalar; >= 1 = # vector-shaped distribution parameter (cutpoints, mixture - # weights). At most one vector-typed position is supported and - # it must be the last variable-typed slot; the vector consumes - # every remaining column of the stacked input. + # weights / locations / scales). Any number of vector-typed + # positions is supported: each consumes the number of columns + # recorded in its ``param_spec`` dim, so the stacked input is + # split by offset. The one exception is a vector whose event + # dimension was not resolved at construction (dim recorded as + # 1 while the runtime tensor is wider); a single such vector, + # in the last slot, absorbs the surplus columns. self._param_event_ranks: tuple[int, ...] = ( param_event_ranks if param_event_ranks is not None @@ -220,16 +224,9 @@ def __init__( f"param_spec length {len(param_spec)}" ) vec_positions = [i for i, r in enumerate(self._param_event_ranks) if r >= 1] - if len(vec_positions) > 1: - raise ValueError( - "MixedInlineDistribution: at most one vector-typed " - "parameter position is supported" - ) - if vec_positions and vec_positions[0] != len(param_spec) - 1: - raise ValueError( - "MixedInlineDistribution: vector-typed parameter must be the last slot" - ) self._has_vector_param: bool = bool(vec_positions) + self._last_vector_pos: int | None = vec_positions[-1] if vec_positions else None + self._n_vector_params: int = len(vec_positions) @property def support(self) -> _constraints.Constraint: @@ -254,28 +251,41 @@ def _resolve_params(self, x: torch.Tensor) -> list[torch.Tensor]: """ if x.dim() == 1: x = x.unsqueeze(-1) - params = [] + params: list[torch.Tensor] = [] var_offset = 0 - n_specs = len(self._param_spec) + # Columns the recorded dims account for; any surplus belongs to + # a single trailing vector whose event dim was not resolved at + # construction. With more than one vector param the surplus is + # unattributable, so the dims must be exact. + known_total = sum(int(v) for k, v in self._param_spec if k == "var") + surplus = x.shape[-1] - known_total + if surplus > 0 and self._n_vector_params > 1: + raise ValueError( + "MixedInlineDistribution: cannot resolve parameters: the " + f"stacked width {x.shape[-1]} exceeds the declared " + f"variable dimension {known_total}, but there is more than " + "one vector-typed parameter, so the surplus columns cannot " + "be attributed to a single parameter" + ) for pos, (kind, value) in enumerate(self._param_spec): if kind == "lit": - lit_val = torch.full( - (x.shape[0],), float(value), device=x.device, dtype=x.dtype + params.append( + torch.full( + (x.shape[0],), float(value), device=x.device, dtype=x.dtype + ) ) - params.append(lit_val) continue - is_last_vector = self._has_vector_param and pos == n_specs - 1 - if is_last_vector: - # Consume all remaining columns as the vector param. - params.append(x[..., var_offset:]) - var_offset = x.shape[-1] + dim = int(value) + if surplus > 0 and pos == self._last_vector_pos: + # The single trailing under-counted vector absorbs the + # surplus columns. + dim += surplus + is_vector = self._param_event_ranks[pos] >= 1 + if dim == 1 and not is_vector: + params.append(x[..., var_offset]) else: - dim = int(value) - if dim == 1: - params.append(x[..., var_offset]) - else: - params.append(x[..., var_offset : var_offset + dim]) - var_offset += dim + params.append(x[..., var_offset : var_offset + dim]) + var_offset += dim return params def rsample( diff --git a/src/quivers/continuous/morphisms.py b/src/quivers/continuous/morphisms.py index a4d54aeb..0420b42a 100644 --- a/src/quivers/continuous/morphisms.py +++ b/src/quivers/continuous/morphisms.py @@ -201,104 +201,6 @@ def __repr__(self) -> str: return f"{cls}({self.domain!r} -> {self.codomain!r})" -class _LookupSource(nn.Module): - """Parameter source for discrete domains: index into a table. - - For each domain element i, returns a parameter vector table[i]. - """ - - def __init__(self, n_entries: int, param_dim: int) -> None: - super().__init__() - self.table = nn.Parameter(torch.randn(n_entries, param_dim) * 0.1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Look up parameters. - - Parameters - ---------- - x : torch.Tensor - Integer indices. Shape (batch,). - - Returns - ------- - torch.Tensor - Parameter vectors. Shape (batch, param_dim). - """ - return self.table[x.long()] - - -class _NeuralSource(nn.Module): - """Parameter source for continuous domains: a small MLP. - - Maps continuous inputs to parameter vectors via a two-layer - network with tanh activations. - """ - - def __init__(self, input_dim: int, param_dim: int, hidden_dim: int = 64) -> None: - super().__init__() - self.net = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.Tanh(), - nn.Linear(hidden_dim, hidden_dim), - nn.Tanh(), - nn.Linear(hidden_dim, param_dim), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Compute parameters from continuous input. - - Parameters - ---------- - x : torch.Tensor - Continuous inputs. Shape (batch, input_dim). - - Returns - ------- - torch.Tensor - Parameter vectors. Shape (batch, param_dim). - """ - if x.dim() == 1: - x = x.unsqueeze(-1) - return self.net(x) - - -def _make_source( - domain: AnySpace, - param_dim: int, - hidden_dim: int = 64, - param_source=None, -) -> nn.Module: - """Create an appropriate parameter source for the given domain. - - Parameters - ---------- - domain : SetObject or ContinuousSpace - The domain space. - param_dim : int - Output dimensionality of the parameter source. - hidden_dim : int - Hidden layer width for the default MLP source when - `param_source` is not supplied. - param_source : ParamSource, optional - Explicit parameter source; overrides the default (MLP for - continuous domains, lookup table for discrete). Use this to - drop in a `LinearSource`, `AttentionSource`, or user-defined - `ParamSource` without subclassing the conditional family. - - Returns - ------- - nn.Module - A `ParamSource` (`_LookupSource` / `_NeuralSource` in the - default path; whatever the caller supplies otherwise). - """ - if param_source is not None: - return param_source - if isinstance(domain, SetObject): - return _LookupSource(domain.size, param_dim) - else: - return _NeuralSource(cast(ContinuousSpace, domain).dim, param_dim, hidden_dim) - - class SampledComposition(ContinuousMorphism): """Composition of morphisms via ancestral sampling. diff --git a/src/quivers/continuous/param_source.py b/src/quivers/continuous/param_source.py index 96d68aed..888cd1d9 100644 --- a/src/quivers/continuous/param_source.py +++ b/src/quivers/continuous/param_source.py @@ -10,13 +10,14 @@ The primitives: * [`LinearSource`][quivers.continuous.param_source.LinearSource]: - one `nn.Linear`, no nonlinearity. Matches the single-linear - layer the transpile backends emit exactly, so a kernel morphism - configured with this source is numerically equivalent to its + the default; one `nn.Linear`, no nonlinearity. Matches the + single-linear layer the transpile backends emit exactly, so a + kernel morphism on this source is numerically equivalent to its transpiled counterpart. * [`MLPSource`][quivers.continuous.param_source.MLPSource]: - the default, parameterised with user-configurable hidden widths - and activation. + a multi-layer perceptron with user-configurable hidden widths and + activation. Selected by `[param_source=mlp]`; a kernel is linear + unless it asks for this. * [`LookupSource`][quivers.continuous.param_source.LookupSource]: a learnable per-entry embedding table, the discrete-domain standard. @@ -36,12 +37,17 @@ building sequential architectures out of primitives. The DSL surface accepts a `[param_source=...]` option on -`morphism` declarations: +`morphism` declarations, with the hidden widths given either as the +option's arguments or through `hidden_dim`: - morphism trans : State -> State [role=kernel, param_source=linear] ~ Normal - morphism trans : State -> State [role=kernel, param_source=mlp(64, 64)] ~ Normal + morphism trans : State -> State [param_source=linear] ~ Normal + morphism trans : State -> State [param_source=mlp] ~ Normal + morphism trans : State -> State [param_source=mlp(64, 64)] ~ Normal + morphism trans : State -> State [param_source=mlp, hidden_dim=[64, 32]] ~ Normal -with the string form parsed by +One width per hidden layer, so the sequence says how many as well as +how wide; a bare `hidden_dim=64` is the one-layer case. The call form +is parsed by [`param_source_from_option`][quivers.continuous.param_source.param_source_from_option]. """ @@ -315,23 +321,62 @@ def forward(self, x: Tensor) -> Tensor: return self.outer(self.inner(x)) +#: The kind a family gets when nothing selects one. +_DEFAULT_SOURCE_KIND = "linear" + + +def _split_source_option(option_value: str | None) -> tuple[str, dict]: + """Split a ``[param_source=...]`` option into its kind and the + arguments it carries. + + ``"mlp"`` is a bare kind; ``"mlp(64, 32)"`` carries positional + hidden widths; ``"attention(heads=4)"`` carries keywords. Absent, + the kind is the default. + """ + if option_value is None: + return _DEFAULT_SOURCE_KIND, {} + val = option_value.strip() + if "(" not in val: + return val, {} + kind, rest = val.split("(", 1) + rest = rest.rstrip(")").strip() + if not rest: + return kind.strip(), {} + kwargs: dict[str, object] = {} + positional: list[int] = [] + for token in rest.split(","): + token = token.strip() + if "=" in token: + k, v = token.split("=", 1) + kwargs[k.strip()] = int(v.strip()) + else: + positional.append(int(token)) + if positional: + kwargs["hidden_dims"] = tuple(positional) + return kind.strip(), kwargs + + def make_param_source( domain: AnySpace, param_dim: int, - kind: str = "mlp", + kind: str = _DEFAULT_SOURCE_KIND, **kwargs, ) -> ParamSource: """Factory that dispatches the `[param_source=...]` DSL option - to the concrete class. Existing families call this to preserve - the pre-abstraction default (MLP with `hidden_dim=64`) when no - option is supplied. + to the concrete class. + + The default is `LinearSource`, so a morphism declared + ``f : X -> Y ~ Normal`` maps its input to the family's parameters + the way its arrow reads: linearly. A model that wants a + nonlinearity between its input and its parameters asks for one, + and the fact then appears in the source rather than in a default. Recognised kinds: * ``"lookup"`` — always used when the domain is a `SetObject`, regardless of the requested kind. - * ``"linear"`` — one `nn.Linear`. - * ``"mlp"`` — the default; `hidden_dims=(64, 64)` unless - overridden by ``hidden_dims`` or ``hidden_dim`` kwargs. + * ``"linear"`` — the default; one `nn.Linear`. + * ``"mlp"`` — `hidden_dims=(64, 64)` unless overridden by + ``hidden_dims`` or ``hidden_dim`` kwargs. * ``"identity"`` — pass through. * ``"attention"`` — self-attention head. """ @@ -359,6 +404,52 @@ def make_param_source( raise ValueError(f"make_param_source: unknown kind {kind!r}") +#: Source kinds with hidden layers, so a width is something they read. +_HIDDEN_LAYER_KINDS: frozenset[str] = frozenset({"mlp"}) + + +def _make_source( + domain: AnySpace, + param_dim: int, + hidden_dim: int | Sequence[int] | None = None, + param_source: ParamSource | None = None, + param_source_option: str | None = None, +) -> ParamSource: + """Create the parameter source a conditional family reads from. + + This is the one seam every family builds its source through, so + each of them takes the same three ways of choosing it: + + * ``param_source``: an instance, which wins outright. + * ``param_source_option``: the DSL's ``[param_source=...]`` text, + a kind with optional arguments (``"mlp"``, ``"mlp(64, 32)"``). + * ``hidden_dim``: the hidden widths, as one width or a sequence of + them. A sequence gives one hidden layer per entry. + + Without any of them the default is linear for a continuous domain + and a lookup table for a discrete one. + + A width that the chosen source has no layers to apply it to raises + rather than being dropped: asking a single matrix how wide its + hidden layers should be has no answer, and silence would read as + one. + """ + if param_source is not None: + return param_source + kind, kwargs = _split_source_option(param_source_option) + if hidden_dim is not None: + if kind not in _HIDDEN_LAYER_KINDS: + raise ValueError( + f"hidden_dim is not read by the {kind!r} parameter source, " + f"which has no hidden layers; select one that does with " + f"param_source, or drop the width" + ) + if "hidden_dims" not in kwargs: + widths = (hidden_dim,) if isinstance(hidden_dim, int) else tuple(hidden_dim) + kwargs["hidden_dims"] = tuple(int(w) for w in widths) + return make_param_source(domain, param_dim, kind=kind, **kwargs) + + def param_source_from_option( domain: AnySpace, param_dim: int, @@ -375,26 +466,7 @@ def param_source_from_option( Unrecognised syntax raises `ValueError` so parse errors surface at compile time rather than as silent identity fallthrough. """ - if option_value is None: - return make_param_source(domain, param_dim) - val = option_value.strip() - if "(" not in val: - return make_param_source(domain, param_dim, kind=val) - kind, rest = val.split("(", 1) - rest = rest.rstrip(")").strip() - if not rest: - return make_param_source(domain, param_dim, kind=kind) - kwargs: dict[str, object] = {} - positional: list[int] = [] - for token in rest.split(","): - token = token.strip() - if "=" in token: - k, v = token.split("=", 1) - kwargs[k.strip()] = int(v.strip()) - else: - positional.append(int(token)) - if positional: - kwargs["hidden_dims"] = tuple(positional) + kind, kwargs = _split_source_option(option_value) return make_param_source(domain, param_dim, kind=kind, **kwargs) diff --git a/src/quivers/continuous/plate.py b/src/quivers/continuous/plate.py index 808eef57..54b89b87 100644 --- a/src/quivers/continuous/plate.py +++ b/src/quivers/continuous/plate.py @@ -291,8 +291,10 @@ class VectorisedObserve(ContinuousMorphism): family : ContinuousMorphism The per-observation distribution family. response : torch.Tensor - Observed values ``r_obs`` of shape ``(N, *codom.shape)`` - (or ``(N,)`` for scalar codomains). + Observed values ``r_obs`` of shape ``(N, *codom.shape)``. A + one-dimensional continuous codomain also accepts ``(N,)``, + which gains its event axis in `_align_response`; a discrete + codomain takes ``(N,)`` indices and has no event axis. """ def __init__(self, family: ContinuousMorphism, response: torch.Tensor) -> None: @@ -318,6 +320,15 @@ def _param_event_ranks(self) -> tuple[int, ...] | None: """ return getattr(self._family, "_param_event_ranks", None) + @property + def _param_spec(self) -> list[tuple[str, int | float]] | None: + """Forward the per-parameter spec (kinds and dims) from the + wrapped family so `_resolve_input` can stack multiple vector- + typed parameters (e.g. `MixtureNormal` weights / locs / scales) + by their declared dims rather than by tensor rank. + """ + return getattr(self._family, "_param_spec", None) + def rsample( self, x: torch.Tensor, sample_shape: torch.Size = torch.Size() ) -> torch.Tensor: @@ -330,6 +341,47 @@ def rsample( """ return self._family.rsample(x, sample_shape) + def _align_response(self, target: torch.Tensor) -> torch.Tensor: + """Give the response its codomain's event axis before scoring. + + A conditional family reads its parameters out of a network, so + for a codomain of dimension ``d`` they arrive as ``(N, d)``. The + natural way to write N scalar responses against a ``d = 1`` + codomain is ``(N,)``, and subtracting that from an ``(N, 1)`` + mean broadcasts to ``(N, N)``. Because the score is then summed, + the result is a finite number that is simply wrong, so the axis + is restored here. + + An inline family is the other layout and is left alone: its + parameters come from the program's environment as per-row + scalars already aligned with the plate, so the response's last + axis indexes the plate rather than an event. ``_param_spec`` + marks that case. + """ + if getattr(self._family, "_param_spec", None) is not None: + return target + codomain = self._family.codomain + dim = getattr(codomain, "dim", None) + if dim is None: + # Discrete codomain: one index per row, no event axis. + return target + dim = int(dim) + if target.dim() == 1: + if dim != 1: + raise ValueError( + f"observed response for codomain {codomain!r} has shape " + f"{tuple(target.shape)}; expected (N, {dim}) for a " + f"{dim}-dimensional codomain" + ) + return target.unsqueeze(-1) + if target.shape[-1] != dim: + raise ValueError( + f"observed response for codomain {codomain!r} has shape " + f"{tuple(target.shape)}; its last axis must be the " + f"codomain's event axis of size {dim}" + ) + return target + def log_prob(self, x: torch.Tensor, y: torch.Tensor | None = None) -> torch.Tensor: """Sum of per-observation log-densities. @@ -338,7 +390,7 @@ def log_prob(self, x: torch.Tensor, y: torch.Tensor | None = None) -> torch.Tens for fast prior-predictive checks. """ target = y if y is not None else cast("torch.Tensor", self._response) - return self._family.log_prob(x, target).sum() + return self._family.log_prob(x, self._align_response(target)).sum() def log_likelihood(self, theta: torch.Tensor) -> torch.Tensor: """Alias for ``log_prob(theta)``; preserved for the Python diff --git a/src/quivers/continuous/programs.py b/src/quivers/continuous/programs.py index de9a9d1e..84b11191 100644 --- a/src/quivers/continuous/programs.py +++ b/src/quivers/continuous/programs.py @@ -389,48 +389,76 @@ def _resolve_input( if len(spec.args) == 1: return self._promote_rank(_lookup_arg(env, spec.args[0])) - # multiple args: stack along feature dimension + # multiple args: stack along the feature dimension. When the + # morphism declares per-parameter event ranks and dims (an + # inline distribution), the stacking interprets each argument by + # its declared role rather than guessing from tensor rank: a + # rank-0 position is a per-row scalar; a rank->=1 position is a + # vector feature of its declared dim, arriving either as a + # shared ``(D,)`` vector or a per-row ``(N, D)`` block. parts = [_lookup_arg(env, a) for a in spec.args] - # Families that declare a vector-typed final parameter (e.g. - # `OrderedLogistic(predictor, cutpoints)` where cutpoints is - # a shared or per-row vector) may supply the vector as a 1-D - # tensor of shape ``(D,)``. Broadcast such shared vectors to - # per-row ``(batch, D)`` before stacking so `_stack_tensors` - # sees uniform per-row shapes. morph = self._modules.get(spec.morphism_name) event_ranks = getattr(morph, "_param_event_ranks", None) - if event_ranks is not None and len(event_ranks) == len(parts): - parts = self._broadcast_vector_params(parts, event_ranks) + param_spec = getattr(morph, "_param_spec", None) + if event_ranks is not None and param_spec is not None: + var_ranks = tuple( + event_ranks[i] for i, (k, _) in enumerate(param_spec) if k == "var" + ) + var_dims = tuple(int(v) for k, v in param_spec if k == "var") + if len(var_ranks) == len(parts): + return self._stack_params(parts, var_ranks, var_dims) return self._stack_tensors(parts) @staticmethod - def _broadcast_vector_params( + def _stack_params( parts: list[torch.Tensor], event_ranks: tuple[int, ...], - ) -> list[torch.Tensor]: - """Broadcast shared vector-typed parameters to per-row shape. - - For each position with ``event_ranks[i] >= 1`` whose tensor - arrives as 1-D shape ``(D,)``, expand to ``(batch, D)`` where - ``batch`` is the max leading dim among the rank-0 (per-row - scalar) positions. Rank-0 tensors and already-per-row rank-1 - tensors pass through unchanged. + dims: tuple[int, ...], + ) -> torch.Tensor: + """Stack inline-distribution parameters into a single input. + + Each parameter is reshaped to a per-row ``(batch, dim)`` block + using its declared event rank and dim, and the blocks are + concatenated along the feature axis. A shared parameter (no + leading batch: a scalar ``()``/``(1,)`` or a vector ``(dim,)``) + reshapes to a batch-1 block and broadcasts against any per-row + parameters. When every parameter is shared, the result is a + single batch-1 row; the response batch is supplied downstream by + the observed value. """ - scalar_batches = [ - t.shape[0] - for t, r in zip(parts, event_ranks, strict=True) - if r == 0 and t.dim() >= 1 - ] - if not scalar_batches: - return parts - batch = max(scalar_batches) - out: list[torch.Tensor] = [] - for t, r in zip(parts, event_ranks, strict=True): - if r >= 1 and t.dim() == 1 and t.shape[0] != batch: - out.append(t.unsqueeze(0).expand(batch, -1)) + shaped: list[torch.Tensor] = [] + for t, rank, dim in zip(parts, event_ranks, dims, strict=True): + tf = t.float() + if rank >= 1: + # Vector feature of size ``dim``. + if tf.dim() <= 1: + # Shared vector ``(dim,)`` (or a scalar broadcast to + # the vector width): one row. + shaped.append(tf.reshape(1, -1)) + else: + shaped.append(tf.reshape(tf.shape[0], -1)) + else: + # Per-row scalar. + if tf.dim() == 0: + shaped.append(tf.reshape(1, 1)) + elif tf.dim() == 1: + shaped.append(tf.unsqueeze(-1)) + else: + shaped.append(tf.reshape(tf.shape[0], -1)) + batch = max((t.shape[0] for t in shaped), default=1) + broadcast: list[torch.Tensor] = [] + for t in shaped: + if t.shape[0] == batch: + broadcast.append(t) + elif t.shape[0] == 1: + broadcast.append(t.expand(batch, *t.shape[1:])) else: - out.append(t) - return out + raise RuntimeError( + "_stack_params: incompatible batch sizes " + f"{[t.shape[0] for t in shaped]}; expected each to be " + f"1 (shared) or {batch} (per-row)" + ) + return torch.cat(broadcast, dim=-1) @staticmethod def _promote_rank(t: torch.Tensor) -> torch.Tensor: diff --git a/src/quivers/dsl/compiler/_options.py b/src/quivers/dsl/compiler/_options.py index 10b30393..cda61092 100644 --- a/src/quivers/dsl/compiler/_options.py +++ b/src/quivers/dsl/compiler/_options.py @@ -158,6 +158,111 @@ def get_option_string( return entry.value.value +def _render_option_value(value: OptionValue) -> str: + """Render an option value back to its surface text.""" + if isinstance(value, OptionName | OptionString): + return value.value + if isinstance(value, OptionNumber): + f = float(value.value) + return str(int(f)) if f.is_integer() else str(f) + if isinstance(value, OptionCall): + args = ", ".join(_render_option_value(a) for a in value.args) + return f"{value.func}({args})" + raise CompileError( + f"option value of kind {type(value).__name__!r} has no surface text", + 0, + 0, + ) + + +def get_option_int_list( + options: tuple[OptionEntry, ...], + key: str, + *, + line: int = 0, + col: int = 0, + default: tuple[int, ...] = (), +) -> tuple[int, ...]: + """Decode a list-of-integers option (``hidden_dim=[64, 32]``). + + Accepts two surface shapes, mirroring `get_option_name_list`: + + * ``[hidden_dim=[64, 32]]`` -> OptionList of OptionNumbers. + * ``[hidden_dim=64]`` -> single OptionNumber, lifted to ``(64,)``. + + One entry per layer, so a sequence says how many as well as how + wide, which is what an MLP needs and a single number cannot say. + """ + entry = find_option(options, key) + if entry is None: + return default + v = entry.value + + def _as_int(value: OptionValue) -> int: + if not isinstance(value, OptionNumber): + ln, cl = _at(line, col, entry) + raise CompileError( + f"option {key!r}: expected an integer or a list of them, " + f"got {type(value).__name__}", + ln, + cl, + ) + number = float(value.value) + if not number.is_integer(): + ln, cl = _at(line, col, entry) + raise CompileError( + f"option {key!r}: expected whole numbers, got {number}", + ln, + cl, + ) + return int(number) + + if isinstance(v, OptionNumber): + return (_as_int(v),) + if isinstance(v, OptionList): + return tuple(_as_int(item) for item in v.items) + ln, cl = _at(line, col, entry) + raise CompileError( + f"option {key!r}: expected an integer or a list of them, such as " + f"``{key}=64`` or ``{key}=[64, 32]``, got {type(v).__name__}", + ln, + cl, + ) + + +def get_option_call_text( + options: tuple[OptionEntry, ...], + key: str, + *, + line: int = 0, + col: int = 0, + default: str | None = None, +) -> str | None: + """Decode an option written as a bare name or a call, to its + surface text. + + Some options name a construction whose arguments are part of the + choice rather than separate keys: ``param_source=mlp`` and + ``param_source=mlp(64, 64)`` select the same architecture at + different widths. The grammar admits both, so both decode here, + to ``"mlp"`` and ``"mlp(64, 64)"``, and the consumer parses the + arguments it defines. + """ + entry = find_option(options, key) + if entry is None: + return default + if not isinstance(entry.value, OptionName | OptionString | OptionCall): + ln, cl = _at(line, col, entry) + raise CompileError( + f"option {key!r}: expected a name or a call such as " + f"``{key}=mlp`` or ``{key}=mlp(64, 64)``, got " + f"{type(entry.value).__name__}", + ln, + cl, + ) + return _render_option_value(entry.value) + + @overload def get_option_int( options: tuple[OptionEntry, ...], diff --git a/src/quivers/dsl/compiler/declarations.py b/src/quivers/dsl/compiler/declarations.py index 5c06933a..52fbf42b 100644 --- a/src/quivers/dsl/compiler/declarations.py +++ b/src/quivers/dsl/compiler/declarations.py @@ -9,6 +9,7 @@ from __future__ import annotations +import inspect import math from collections.abc import Callable @@ -64,11 +65,12 @@ from quivers.dsl.compiler._options import ( check_option_keys, find_option, + get_option_call_text, get_option_float, + get_option_int_list, get_option_int, get_option_name, get_option_name_list, - get_option_string, ) from quivers.dsl.compiler._prelude import ( _ALGEBRA_REGISTRY, @@ -123,6 +125,44 @@ } ) +# The union above spans every role, so checking against it alone would +# accept a key the chosen lowering never reads and drop it in silence. +# Each role therefore declares the keys it actually consumes, and a key +# outside its set is rejected rather than ignored: `scale` configures a +# latent morphism's init and means nothing to a family-backed kernel, +# whose parameters come from its `param_source`. +# +# ``role`` and ``replicate`` pick the lowering and its multiplicity, so +# every role reads them. +_COMMON_MORPHISM_OPTION_KEYS: frozenset[str] = frozenset({"role", "replicate"}) + +_ROLE_OPTION_KEYS: dict[str, frozenset[str]] = { + # `_compile_latent_role` reads scale / init; a `~ Family(...)` init + # carries its axis-role clause through `_validate_family_axes`. + "latent": frozenset({"scale", "init", "over", "iid"}), + "observed": frozenset(), + # `_make_continuous_morphism` threads these into the family. + "kernel": frozenset( + { + "n_layers", + "hidden_dim", + "param_source", + "rank", + "temperature", + "over", + "iid", + } + ), + "embed": frozenset(), + "discretize": frozenset({"bins"}), + "let": frozenset({"over", "iid"}), +} + + +def _signature_params(cls) -> frozenset[str]: + """The keyword names a family's ``__init__`` accepts.""" + return frozenset(inspect.signature(cls.__init__).parameters) + def _apply_auto_init(morph, domain, codomain, algebra) -> None: """Apply the algebra's saturation-free init recipe to a freshly @@ -672,6 +712,7 @@ def _compile_morphism_named(self, decl: MorphismDecl, name: str) -> None: decl.col, ) role = self._resolve_morphism_role(decl, name) + self._check_role_options(decl, name, role) replicate = get_option_int( decl.options, "replicate", @@ -698,6 +739,35 @@ def _compile_morphism_named(self, decl: MorphismDecl, name: str) -> None: if replicate is not None: self._groups[name] = names + def _check_role_options(self, decl: MorphismDecl, name: str, role: str) -> None: + """Reject an option the resolved role's lowering never reads. + + `check_option_keys` has already rejected keys outside the + union, so anything reaching here is a real morphism option + applied to a role that ignores it. Ignoring it silently is how + ``[scale=0.5] ~ Normal`` came to read as a configured init + while doing nothing at all. + """ + allowed = _COMMON_MORPHISM_OPTION_KEYS | _ROLE_OPTION_KEYS[role] + for entry in decl.options: + if entry.key in allowed: + continue + owners = sorted( + r for r, keys in _ROLE_OPTION_KEYS.items() if entry.key in keys + ) + where = ( + f"it configures {', '.join(f'role={r}' for r in owners)}" + if owners + else "no role reads it" + ) + raise CompileError( + f"morphism {name!r}: option {entry.key!r} is not read by " + f"role={role}; {where}. Remove it, or declare the " + f"morphism under a role that reads it.", + entry.line or decl.line, + entry.col or decl.col, + ) + def _resolve_morphism_role(self, decl: MorphismDecl, name: str) -> str: """Resolve a morphism's role: explicit ``role=`` wins; absent role defaults to ``kernel``. @@ -1059,36 +1129,42 @@ def _make_continuous_morphism( decl.col, ) cls = registry[family_name] - hidden_dim = get_option_int( + hidden_dim = get_option_int_list( decl.options, "hidden_dim", line=decl.line, col=decl.col, - default=64, ) - kwargs: dict = {"hidden_dim": int(hidden_dim)} - # Optional `[param_source=[(...)]]` DSL surface for - # picking the parameter-source architecture (linear, MLP, - # attention, identity). The default MLP with `hidden_dim` - # matches the pre-abstraction behaviour. The kwarg is - # threaded through to the conditional family's `__init__`, - # which uses `param_source_from_option` internally to build - # the concrete `ParamSource` once `param_dim` is knowable. - param_source_opt = get_option_name( + kwargs: dict = {} + if hidden_dim: + kwargs["hidden_dim"] = hidden_dim[0] if len(hidden_dim) == 1 else hidden_dim + # Optional `[param_source=]` / `[param_source=(...)]` + # DSL surface for picking the parameter-source architecture + # (linear, MLP, attention, identity). The default is linear. + # Both keys are threaded to the family's `__init__`, which hands + # them to `_make_source`: the option's text carries any + # parenthesised widths, so the call form has to reach it as + # surface text rather than as a bare name. + param_source_opt = get_option_call_text( decl.options, "param_source", line=decl.line, col=decl.col, ) - if param_source_opt is None: - param_source_opt = get_option_string( - decl.options, - "param_source", - line=decl.line, - col=decl.col, - ) if param_source_opt is not None: kwargs["param_source_option"] = param_source_opt + # A family whose parameters do not come from a source has + # nothing to point these at. Saying so beats a TypeError out of + # a constructor with no line or column on it. + unread = sorted(k for k in kwargs if k not in _signature_params(cls)) + if unread: + raise CompileError( + f"morphism {decl.names[0]!r}: {family_name} does not take " + f"{', '.join(repr(k) for k in unread)}; its parameters do " + f"not come from a ``param_source``", + decl.line, + decl.col, + ) rank = get_option_int( decl.options, "rank", diff --git a/src/quivers/dsl/compiler/programs.py b/src/quivers/dsl/compiler/programs.py index e5ceb9b5..ac0f005b 100644 --- a/src/quivers/dsl/compiler/programs.py +++ b/src/quivers/dsl/compiler/programs.py @@ -92,6 +92,29 @@ ) +def _fibration_index(env: dict[str, torch.Tensor], name: str) -> torch.Tensor: + """Resolve a grouped-marginalize ``via`` fibration index from the + runtime environment. + + The index is a per-row map from the response plate into the + grouping plate, supplied as host data through the observations + dict (like any free covariate). A name that resolves to nothing + is a user error, so it fails with a clear message rather than a + bare ``KeyError``. + """ + idx = env.get(name) + if idx is None: + raise CompileError( + f"grouped marginalize: the ``via`` fibration index {name!r} " + f"is neither a bound program variable nor supplied at runtime; " + f"pass it in the observations dict as a long tensor of " + f"row-to-group indices", + 0, + 0, + ) + return idx + + # Value carried by a let-binding at compile time. The let # sublanguage is a small typed lambda calculus over heterogeneous # values: numeric literals lift to torch tensors; identifier @@ -2437,13 +2460,13 @@ def _marginalize_grouped_callable( if fib_axes is not None: idx_tuple = [] for axis_name in fib_axes: - idx = env[axis_name] + idx = _fibration_index(env, axis_name) if idx.dim() == 2 and idx.shape[-1] == 1: idx = idx.squeeze(-1) idx_tuple.append(idx.to(torch.long)) idx_list.append(tuple(idx_tuple)) elif fib_var is not None: - idx = env[fib_var] + idx = _fibration_index(env, fib_var) if idx.dim() == 2 and idx.shape[-1] == 1: idx = idx.squeeze(-1) idx_list.append(idx.to(torch.long)) diff --git a/src/quivers/formulas/compile.py b/src/quivers/formulas/compile.py index e76994c0..e0e22c6d 100644 --- a/src/quivers/formulas/compile.py +++ b/src/quivers/formulas/compile.py @@ -30,9 +30,13 @@ * One ``object G : K`` declaration per random-effect grouping factor (with ``K`` levels). * For each fixed-design column ``c``: one scalar latent draw - inside the program body, ``beta_c <- Normal(0, fixed_prior)``. - The per-row covariate values for ``c`` flow in as a free - variable via the host-data channel (``observations[c]``). + inside the program body, ``beta_c <- Normal(0, fixed_prior)``, + whose scale is autoscaled by ``1 / rms(c)`` so the nominal + value reads in contribution space and means the same thing on + a raw predictor and on an orthonormal ``poly`` column. An + explicit per-name prior is emitted as written. The per-row + covariate values for ``c`` flow in as a free variable via the + host-data channel (``observations[c]``). * For each random-effect group ``(slope | g)``: a `HalfNormal` scale latent plus a per-level plate draw, with the per-row contribution as a plate-gather @@ -87,6 +91,51 @@ ) +#: Prior families whose last positional argument is a scale, so an +#: autoscaled default divides it by the column's RMS. +_SCALE_LAST_PRIOR_FAMILIES: frozenset[str] = frozenset( + {"Normal", "Cauchy", "Laplace", "StudentT", "Logistic"} +) + + +def _autoscaled_prior_args( + family_name: str, + args: tuple[str | float, ...], + column: "np.ndarray", +) -> tuple[str | float, ...]: + """Rescale a default coefficient prior to the column it multiplies. + + A prior stated on the coefficient alone is a statement about the + coefficient's *contribution*, which is what the response actually + sees: the column enters the linear predictor as ``beta * column``, + so ``beta ~ Normal(0, s)`` asserts a contribution of scale + ``s * rms(column)``. The same nominal prior therefore means + something different for every column. + + That bites hardest on an orthonormal basis. ``poly(x, k)`` returns + columns of norm one, whose entries are of order ``1 / sqrt(N)``, so + a fixed ``Normal(0, 5)`` asserts a contribution near zero and the + fit obliges: the noise scale rises to the marginal spread of the + response and the coefficients never leave the prior. + + Dividing the scale by the column's RMS states the prior in + contribution space instead, which is what the nominal value reads + as and what makes it mean the same thing across columns. The + coefficients stay on their own column's scale, so nothing has to be + transformed back afterwards. This is the autoscaling convention + `rstanarm` applies by default. + """ + if family_name not in _SCALE_LAST_PRIOR_FAMILIES or not args: + return args + scale = args[-1] + if not isinstance(scale, float): + return args + rms = float(np.sqrt(np.mean(np.square(np.asarray(column, dtype=np.float64))))) + if not np.isfinite(rms) or rms <= 0.0: + return args + return (*args[:-1], scale / rms) + + def _parse_prior_call(text: str) -> tuple[str, tuple[str | float, ...]]: """Split a brms-style prior template ``"Family(arg, arg, ...)"`` into its family name and a tuple of argument tokens. Numeric @@ -528,8 +577,16 @@ def _build_module(self, formula: Formula) -> Module: # Fixed effects: one scalar latent per design-matrix column. for col in formula.fixed_columns: beta_name = "intercept" if col.is_intercept else f"beta_{col.qvr_name}" - prior_text = self._user_priors.get(beta_name, self._fixed_prior) + override = self._user_priors.get(beta_name) + prior_text = override if override is not None else self._fixed_prior family_name, args = _parse_prior_call(prior_text) + # A default prior is autoscaled to the column it multiplies, + # so its nominal scale reads in contribution space and means + # the same thing on a raw predictor and on an orthonormal + # `poly` column. An explicit prior is the user's statement + # about that coefficient and is emitted as written. + if override is None and not col.is_intercept: + args = _autoscaled_prior_args(family_name, args, col.data) program_steps.append(_draw(beta_name, family_name, args)) if col.is_intercept: linear_terms.append(_var(beta_name)) diff --git a/tests/test_deduction.py b/tests/test_deduction.py index 1e5b6b01..40299552 100644 --- a/tests/test_deduction.py +++ b/tests/test_deduction.py @@ -309,3 +309,56 @@ def test_nuts_runs_and_log_density_does_not_collapse(): assert int(res.divergence_counts.sum()) == 0, ( f"NUTS reported divergences: {res.divergence_counts.tolist()}" ) + + +def test_lambda_body_occurrence_matches_its_binder(): + """A lexicon LF's bound occurrence must carry its binder's + canonical symbol. + + ``binders Lam`` alpha-renames a lambda's bound variable to a fresh + canonical name per term construction, and the occurrences in the + body have to be renamed with it. A rewrite that renames the binder + while leaving the body's ``Var(x)`` behind still produces a + well-formed tuple and a chart that parses, so nothing downstream + fails: the lambda simply binds nothing and the body references a + free variable. This pins the two together. + """ + from pathlib import Path + + from quivers.dsl import load as _load + + src = Path("docs/examples/source/montague_nli.qvr") + if not src.exists(): + return # docs not vendored in this checkout + ded = _load(str(src)).deductions["Montague"] + chart = ded(["every", "dog", "barks"]) + + # "dog" is Lam(x, App(dog_p, Var(x))): one binder, one occurrence. + lfs = [ + item[4] + for item, _ in chart.chart.items() + if isinstance(item, tuple) and item[:3] == ("span", 1, 2) and len(item) > 4 + ] + assert lfs, "the noun 'dog' did not derive a span(1, 2) item" + lf = lfs[0] + + assert lf[0] == "Lam", f"expected a Lam-headed LF, got {lf!r}" + bound = lf[1] + assert isinstance(bound, tuple) and len(bound) == 1, ( + f"binder should carry one canonical symbol, got {bound!r}" + ) + + def _occurrences(term) -> list: + if not isinstance(term, tuple): + return [] + if term and term[0] == "Var": + return [term[1]] + return [o for sub in term for o in _occurrences(sub)] + + occurrences = _occurrences(lf[2]) + assert occurrences, f"no Var occurrence in the lambda body: {lf!r}" + for occurrence in occurrences: + assert occurrence == bound, ( + f"lambda binds {bound!r} but its body references {occurrence!r}; " + f"the occurrence is unbound" + ) diff --git a/tests/test_dsl_diagnostics.py b/tests/test_dsl_diagnostics.py index 9d463e8c..6966ee07 100644 --- a/tests/test_dsl_diagnostics.py +++ b/tests/test_dsl_diagnostics.py @@ -152,6 +152,28 @@ "'observed'", ), ), + ( + # `scale` configures a latent morphism's init. A family-backed + # kernel draws its parameters from its param_source, so the + # option means nothing there and must not be swallowed. + "scale-on-kernel-names-the-role-that-reads-it", + ("object S : Real 4\nmorphism f : S -> S [scale=0.1] ~ Normal\nexport f\n"), + ("line 2", "'scale'", "not read by role=kernel", "role=latent"), + ), + ( + "bins-on-kernel-names-the-role-that-reads-it", + ("object S : Real 4\nmorphism f : S -> S [bins=4] ~ Normal\nexport f\n"), + ("line 2", "'bins'", "not read by role=kernel", "role=discretize"), + ), + ( + "param-source-on-latent-names-the-role-that-reads-it", + ( + "object A : FinSet 3\n" + "morphism f : A -> A [role=latent, param_source=mlp]\n" + "export f\n" + ), + ("line 2", "'param_source'", "not read by role=latent", "role=kernel"), + ), ( "typo-keyword", ("object A : FinSet 3\nmorphsim f : A -> A\nexport f\n"), diff --git a/tests/test_formulas.py b/tests/test_formulas.py index 6bde7405..be09b159 100644 --- a/tests/test_formulas.py +++ b/tests/test_formulas.py @@ -39,6 +39,8 @@ import numpy as np import pandas as pd import polars as pl +import re + import pytest from quivers.dsl import Compiler, loads @@ -739,12 +741,28 @@ def test_random_intercept_partial_pooling(self): seed=0, ) means = _posterior_means(result.posterior, n_obs=len(ys)) - # Grand mean within tolerance of the true value. - assert abs(means["intercept"].item() - true_grand_mean) < 0.5 - # Per-group random effects are an 8-vector; correlation with - # the true effects should be strong. - post_group = means["alpha_g"].detach().numpy().reshape(-1) + # The formula compiler emits the non-centred parameterisation, + # so a group's random effect is its standard-normal draw scaled + # by the group-level sigma; there is no centred `alpha_g` site. + post_group = ( + (means["sigma_g_Intercept"] * means["z_g_Intercept"]) + .detach() + .numpy() + .reshape(-1) + ) assert len(post_group) == n_groups + + # The intercept and the random effects are identified only up to + # a constant shared shift: adding c to the intercept and taking + # c off every group effect leaves the likelihood untouched, and + # only the N(0, 1) prior on the draws pulls the split back, which + # with eight groups it does weakly. So the level is asserted on + # the identified sum rather than on the intercept alone. + level = means["intercept"].item() + float(post_group.mean()) + assert abs(level - (true_grand_mean + group_effects.mean())) < 0.5 + + # The effects themselves are identified up to that same shift, + # which correlation is invariant to. corr = np.corrcoef(post_group, group_effects)[0, 1] assert corr > 0.7 @@ -767,11 +785,20 @@ def test_polynomial_orthogonal_recovers_quadratic(self): seed=0, ) means = _posterior_means(result.posterior, n_obs=N) - # The orthogonal polynomial coefficients aren't the raw - # quadratic ones, but BOTH should be substantially nonzero - # because the true function has linear + quadratic content. - assert abs(means["beta_poly_x_2_1"].item()) > 0.5 - assert abs(means["beta_poly_x_2_2"].item()) > 0.5 + # `poly` returns an orthonormal basis, so least squares on it is + # just the projection of y onto each column and gives the exact + # target the fit has to reach. Both coefficients are large + # because the columns are of norm one over N rows. + from formulae import design_matrices + + X = np.asarray(design_matrices("y ~ poly(x, 2)", df).common.design_matrix) + exact = np.linalg.lstsq(X, y, rcond=None)[0] + assert abs(means["beta_poly_x_2_1"].item() - exact[1]) < 0.5 * abs(exact[1]) + assert abs(means["beta_poly_x_2_2"].item() - exact[2]) < 0.5 * abs(exact[2]) + # The noise scale is the data's, not the marginal spread of y: + # a fit that gives up puts sigma near std(y) and leaves the + # coefficients at their prior. + assert abs(means["sigma"].item() - 0.3) < 0.15 def test_log_transform_recovers_slope(self): """y ~ log(w) recovers a slope when y is generated from a @@ -892,9 +919,16 @@ def test_random_slope_recovers_per_group_effect(self): seed=0, ) means = _posterior_means(result.posterior, n_obs=len(ys)) - # Per-group random intercepts and slopes correlate with truth. - post_int = means["alpha_g"].detach().numpy().reshape(-1) - post_slope = means["beta_g_x"].detach().numpy().reshape(-1) + # The formula compiler emits the non-centred parameterisation, + # so a group's random effect is its standard-normal draw scaled + # by the group-level sigma; there is no centred `alpha_g` site. + post_int = ( + (means["sigma_g_Intercept"] * means["z_g_Intercept"]) + .detach() + .numpy() + .reshape(-1) + ) + post_slope = (means["sigma_g_x"] * means["z_g_x"]).detach().numpy().reshape(-1) assert np.corrcoef(post_int, ranef_intercept)[0, 1] > 0.7 assert np.corrcoef(post_slope, ranef_slope)[0, 1] > 0.5 @@ -930,3 +964,57 @@ def test_negbin_carries_disp(self, count_df): def test_unknown_family_raises(self, base_df): with pytest.raises(ValueError, match="unknown family"): formula_to_qvr("y ~ x", data=base_df, family="nonexistent") + + +# --------------------------------------------------------------------------- +# Default coefficient priors are autoscaled to their column +# --------------------------------------------------------------------------- + + +class TestPriorAutoscaling: + def test_default_prior_scales_by_the_column_rms(self, base_df): + """A column enters as ``beta * column``, so the default prior's + scale is divided by the column's RMS: the nominal value then + reads in contribution space and means the same thing whatever + the column's units.""" + src = formula_to_qvr("y ~ x", data=base_df, family="gaussian") + rms = float(np.sqrt(np.mean(np.square(base_df["x"].to_numpy())))) + match = re.search(r"sample beta_x <- Normal\(0\.0, ([0-9.eE+-]+)\)", src) + assert match, f"no autoscaled beta_x prior in:\n{src}" + assert float(match.group(1)) == pytest.approx(5.0 / rms, rel=1e-6) + + def test_orthonormal_poly_columns_get_a_wide_prior(self): + """`poly` returns columns of norm one, whose entries run about + 1/sqrt(N). Without autoscaling the default prior would assert + the contribution is near zero and the fit would believe it.""" + rng = np.random.default_rng(0) + df = pd.DataFrame({"y": rng.normal(size=200), "x": rng.uniform(-2, 2, 200)}) + src = formula_to_qvr("y ~ poly(x, 2)", data=df, family="gaussian") + scales = [ + float(s) + for s in re.findall( + r"sample beta_poly_x_2_\d <- Normal\(0\.0, ([0-9.eE+-]+)\)", src + ) + ] + assert len(scales) == 2 + # rms of a norm-one column over N rows is 1/sqrt(N), so the + # scale lands near 5 * sqrt(N). + for s in scales: + assert s == pytest.approx(5.0 * np.sqrt(200), rel=1e-3) + + def test_intercept_prior_is_left_alone(self, base_df): + """The intercept multiplies a column of ones, so there is no + scale to correct for.""" + src = formula_to_qvr("y ~ x", data=base_df, family="gaussian") + assert "sample intercept <- Normal(0.0, 5.0)" in src + + def test_explicit_prior_is_emitted_as_written(self, base_df): + """An override is the user's statement about that coefficient, + so it is not rescaled underneath them.""" + src = formula_to_qvr( + "y ~ x", + data=base_df, + family="gaussian", + priors={"beta_x": "Normal(0.0, 0.25)"}, + ) + assert "sample beta_x <- Normal(0.0, 0.25)" in src diff --git a/tests/test_gallery_sweep.py b/tests/test_gallery_sweep.py index 0095637d..d1b99f74 100644 --- a/tests/test_gallery_sweep.py +++ b/tests/test_gallery_sweep.py @@ -254,14 +254,17 @@ def test_gallery_try_it_blocks_execute(doc_name): """Execute every ``## Try it`` Python block in the doc. A block that names a missing helper or breaks under the current compiler fails the suite, keeping the docs honest about what the framework - supports today. Blocks may opt out via a ```` + supports today. The blocks of a page share one namespace and run + in order, so a later block sees the imports and bindings an + earlier one established, exactly as a reader stepping through the + page would. Blocks may opt out via a ```` HTML comment immediately above the fenced block.""" path = Path(f"docs/examples/{doc_name}") blocks = _extract_try_it_blocks(path.read_text()) if not blocks: pytest.skip(f"{doc_name}: no Try-it blocks") + ns = {"__name__": f"_try_it_{path.stem}"} for i, block in enumerate(blocks): - ns = {"__name__": f"_try_it_{path.stem}_{i}"} try: exec(compile(block, f"{doc_name}::block-{i}", "exec"), ns) except SystemExit: diff --git a/tests/test_grouped_marginalize.py b/tests/test_grouped_marginalize.py index 03f63e17..6e1122ba 100644 --- a/tests/test_grouped_marginalize.py +++ b/tests/test_grouped_marginalize.py @@ -213,7 +213,6 @@ def test_grouped_block_compiles(self): program demo : Item -> Item sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) marginalize cls : Class <- Dirichlet(probs) [over=Item, reduction=logsumexp] observe r : Resp <- HalfNormal(1.0) [via=idx] return probs @@ -256,7 +255,6 @@ def test_grouped_requires_class_annotation(self): program demo : Item -> Item sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) marginalize cls <- Dirichlet(probs) [over=Item] sample inner : Resp <- HalfNormal(1.0) return probs @@ -278,7 +276,6 @@ def test_grouped_undeclared_over_errors(self): program demo : Item -> Item sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) marginalize cls : Class <- Dirichlet(probs) [over=NotAnObject] sample inner : Resp <- HalfNormal(1.0) return probs diff --git a/tests/test_grouped_marginalize_e2e.py b/tests/test_grouped_marginalize_e2e.py index 6c60f879..152fc415 100644 --- a/tests/test_grouped_marginalize_e2e.py +++ b/tests/test_grouped_marginalize_e2e.py @@ -27,9 +27,11 @@ from __future__ import annotations import textwrap +import pytest import torch from quivers.dsl import loads +from quivers.dsl.compiler import CompileError from quivers.inference import ( AutoNormalGuide, ELBO, @@ -55,7 +57,6 @@ def _two_class_mixture_model() -> str: program two_class_mix : Resp -> Resp sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) sample mu_shift <- Normal(0.0, 1.0) marginalize cls : Class <- Dirichlet(probs) [over=Item] observe r : Resp <- Normal(mu_shift, 1.0) [via=idx] @@ -92,6 +93,27 @@ def test_grouped_marginalize_log_joint_returns_finite_scalar() -> None: assert torch.isfinite(out).all() +def test_missing_via_index_raises_clear_error() -> None: + """A ``via`` fibration index is free host data: it must be + supplied through the observations dict at runtime, exactly like + a covariate. When it is neither a bound program variable nor + supplied, ``log_joint`` fails with a clear message that names + the index and says how to supply it, rather than a bare + ``KeyError``.""" + src = _two_class_mixture_model() + model = loads(textwrap.dedent(src)).morphism + obs = { + "probs": torch.tensor([0.6, 0.4]), + "mu_shift": torch.tensor([0.5]), + "_grouped_ll_cls_0": torch.zeros(8, 2), + # ``idx`` intentionally omitted. + } + with pytest.raises(CompileError, match=r"via`` fibration index"): + model.log_joint(torch.zeros(1, 1), obs) + with pytest.raises(CompileError, match="supplied at runtime"): + model.log_joint(torch.zeros(1, 1), obs) + + def test_svi_runs_on_grouped_marginalize_model() -> None: """SVI takes ELBO steps on a model that uses a grouped marginalize block. The continuous latent ``mu_shift`` has a @@ -164,7 +186,6 @@ def test_grouped_marginalize_recovers_mixture_proportions() -> None: program recovery : Resp -> Resp sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) sample mu_shift <- Normal(0.0, 1.0) marginalize cls : Class <- Dirichlet(probs) [over=Item] observe r : Resp <- Normal(mu_shift, 1.0) [via=idx] diff --git a/tests/test_grouped_marginalize_extended.py b/tests/test_grouped_marginalize_extended.py index 8fb58f23..ed69171c 100644 --- a/tests/test_grouped_marginalize_extended.py +++ b/tests/test_grouped_marginalize_extended.py @@ -342,7 +342,6 @@ def test_reduction_clause_compiles(self) -> None: program demo : Resp -> Resp sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) marginalize cls : Class <- Dirichlet(probs) [over=Item, reduction=sum] observe r : Resp <- HalfNormal(1.0) [via=idx] return probs @@ -364,7 +363,6 @@ def test_unknown_reduction_errors(self) -> None: program demo : Resp -> Resp sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) marginalize cls : Class <- Dirichlet(probs) [over=Item, reduction=bogus] observe r : Resp <- HalfNormal(1.0) [via=idx] return probs @@ -398,10 +396,9 @@ def test_categorical_with_literal_first_arg_errors(self) -> None: object Class : FinSet 2 program demo : Resp -> Resp - sample idx : Resp <- HalfNormal(1.0) marginalize cls : Class <- Dirichlet(1.0) [over=Item] observe r : Resp <- HalfNormal(1.0) [via=idx] - return idx + return cls export demo """ with pytest.raises(CompileError, match="named probs"): @@ -423,7 +420,6 @@ def test_body_without_observe_errors(self) -> None: program demo : Resp -> Resp sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) marginalize cls : Class <- Dirichlet(probs) [over=Item] sample other : Resp <- HalfNormal(1.0) return probs diff --git a/tests/test_grouped_marginalize_semantics.py b/tests/test_grouped_marginalize_semantics.py index 8787af66..0888866f 100644 --- a/tests/test_grouped_marginalize_semantics.py +++ b/tests/test_grouped_marginalize_semantics.py @@ -108,7 +108,6 @@ def test_body_with_multiple_lets_using_latent() -> None: program bodylet : Resp -> Resp sample probs : Class <- HalfNormal(1.0) - sample idx : Resp <- HalfNormal(1.0) marginalize cls : Class <- Dirichlet(probs) [over=Item] observe r : Resp <- HalfNormal(1.0) [via=idx] return probs diff --git a/tests/test_inline_multi_vector.py b/tests/test_inline_multi_vector.py new file mode 100644 index 00000000..baf7f465 --- /dev/null +++ b/tests/test_inline_multi_vector.py @@ -0,0 +1,126 @@ +"""Inline distributions with more than one vector-typed parameter. + +``MixedInlineDistribution`` stacks a distribution's variable parameters +into a single input tensor and splits them back out by declared dim, so +a family may declare any number of vector-typed parameters. The +canonical case is ``MixtureNormal(weights, locations, scales)``, whose +three per-component vectors are each shared across the response plate +and define a finite Gaussian mixture scored per row. + +Run with:: + + pytest tests/test_inline_multi_vector.py +""" + +from __future__ import annotations + +import textwrap +from typing import cast + +import torch +import torch.nn as nn + +from quivers.continuous.inline import make_inline_distribution +from quivers.continuous.spaces import Euclidean +from quivers.dsl import loads + + +def _reference_mixture_logprob( + weights: torch.Tensor, + loc: torch.Tensor, + scale: torch.Tensor, + y: torch.Tensor, +) -> torch.Tensor: + """Closed-form per-row log-density of a finite Gaussian mixture.""" + per_component = ( + torch.distributions.Normal(loc, scale).log_prob(y[:, None]) + weights.log() + ) + return torch.logsumexp(per_component, dim=-1) + + +def test_inline_mixture_normal_log_prob_matches_reference() -> None: + """Three shared per-component vectors (weights, loc, scale) split + cleanly and score each row against the closed-form mixture.""" + k, n = 3, 64 + types = { + "w": Euclidean(name="w", dim=k), + "m": Euclidean(name="m", dim=k), + "s": Euclidean(name="s", dim=k), + } + morph, order = make_inline_distribution( + "MixtureNormal", ("w", "m", "s"), Euclidean(name="R", dim=1), types + ) + assert order == ("w", "m", "s") + + weights = torch.tensor([0.3, 0.4, 0.3]) + loc = torch.tensor([-3.0, 0.0, 3.0]) + scale = torch.tensor([0.5, 0.7, 0.4]) + stacked = torch.cat([weights, loc, scale]).reshape(1, -1) + y = torch.randn(n) + + got = morph.log_prob(stacked, y).reshape(-1) + expected = _reference_mixture_logprob(weights, loc, scale, y) + assert got.shape == (n,) + assert torch.allclose(got, expected, atol=1e-5) + + +def test_inline_mixture_normal_per_row_gmm_compiles_and_fits() -> None: + """The per-row Gaussian mixture program compiles and an SVI fit + drives the loss down and recovers the component means.""" + from quivers.inference import ELBO, SVI, AutoNormalGuide + + source = textwrap.dedent( + """ + composition log_prob [level=algebra] + + object Component : FinSet 3 + object Resp : FinSet 300 + + program gmm : Resp -> Resp + sample probs <- Dirichlet(1.0) [over=Component] + sample mu : Component <- Normal(0.0, 5.0) + sample sigma : Component <- HalfNormal(1.0) + observe r : Resp <- MixtureNormal(probs, mu, sigma) + return probs + + export gmm + """ + ) + model = cast(nn.Module, loads(source).morphism) + assert model is not None + + torch.manual_seed(0) + true_probs = torch.tensor([0.3, 0.4, 0.3]) + true_mu = torch.tensor([-4.0, 0.0, 4.0]) + true_sigma = torch.tensor([0.4, 0.5, 0.4]) + comps = torch.distributions.Categorical(true_probs).sample(torch.Size((300,))) + r = torch.distributions.Normal(true_mu[comps], true_sigma[comps]).sample() + x = torch.zeros(300, 1) + obs = {"r": r, "probs": true_probs} + + guide = AutoNormalGuide(model, observed_names={"r", "probs"}) + optim = torch.optim.Adam( + list(model.parameters()) + list(guide.parameters()), lr=3e-2 + ) + svi = SVI(model, guide, optim, ELBO(num_particles=4)) + losses = [svi.step(x, obs) for _ in range(300)] + assert losses[-1] < losses[0] + + +def test_inline_multi_vector_splits_by_declared_dim() -> None: + """Distinct vector dims split at the right offsets: a length-2 and a + length-3 shared vector recover their own slices.""" + types = { + "a": Euclidean(name="a", dim=2), + "b": Euclidean(name="b", dim=3), + } + morph, _ = make_inline_distribution( + "MixtureNormal", ("a", "a", "b"), Euclidean(name="R", dim=1), types + ) + # Two length-2 vectors then a length-3 vector: 7 columns total. + stacked = torch.arange(7.0).reshape(1, -1) + resolved = morph._resolve_params(stacked) + assert [tuple(t.shape) for t in resolved] == [(1, 2), (1, 2), (1, 3)] + assert torch.equal(resolved[0].reshape(-1), torch.tensor([0.0, 1.0])) + assert torch.equal(resolved[1].reshape(-1), torch.tensor([2.0, 3.0])) + assert torch.equal(resolved[2].reshape(-1), torch.tensor([4.0, 5.0, 6.0])) diff --git a/tests/test_param_source.py b/tests/test_param_source.py index eaf8250e..0b3ddcf9 100644 --- a/tests/test_param_source.py +++ b/tests/test_param_source.py @@ -120,9 +120,12 @@ def test_make_param_source_discrete_returns_lookup() -> None: assert isinstance(ps, LookupSource) -def test_make_param_source_continuous_default_mlp() -> None: +def test_make_param_source_continuous_default_linear() -> None: + """A continuous domain defaults to one linear map. The default + decides whether a model is linear, so it is the reading of the + arrow rather than a hidden network.""" ps = make_param_source(Euclidean(name="X", dim=4), 8) - assert isinstance(ps, MLPSource) + assert isinstance(ps, LinearSource) def test_make_param_source_kind_linear() -> None: @@ -156,7 +159,7 @@ def test_param_source_from_option_paren_forms() -> None: def test_param_source_from_option_none_returns_default() -> None: ps = param_source_from_option(Euclidean(name="X", dim=4), 8, None) - assert isinstance(ps, MLPSource) + assert isinstance(ps, LinearSource) def test_conditional_normal_accepts_param_source_kwarg() -> None: @@ -176,13 +179,14 @@ def test_conditional_normal_accepts_param_source_option_string() -> None: assert isinstance(n.param_source, LinearSource) -def test_conditional_normal_default_source_is_neural() -> None: +def test_conditional_normal_default_source_is_linear() -> None: + """A family built without an explicit source goes through the same + factory as the DSL option, so its default is one linear map rather + than a network.""" X = Euclidean(name="X", dim=4) Y = Euclidean(name="Y", dim=2) n = ConditionalNormal(X, Y) - # Default: `_NeuralSource` (pre-abstraction two-layer MLP) or - # `MLPSource` depending on which entry-point built it; both are - # nn.Modules with the right output dim. + assert isinstance(n.param_source, LinearSource) x = torch.randn(3, 4) y = torch.randn(3, 2) assert n.log_prob(x, y).shape == torch.Size([3]) @@ -222,3 +226,177 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Y = Euclidean(name="Y", dim=2) n = ConditionalNormal(X, Y, param_source=src) assert n.param_source is src + + +# --------------------------------------------------------------------------- +# DSL surface: [param_source=...] reaches the family +# --------------------------------------------------------------------------- + + +def _family_of(source: str): + """Compile a one-kernel program and return its conditional family.""" + from quivers.dsl import loads + + prog = loads(source) + model = prog.morphism + assert model is not None + return dict(model.named_modules())["_step_y._family"] + + +def _program(option: str) -> str: + return ( + "object F : Real 1\n" + "object T : Real 1\n" + "object R : FinSet 8\n" + "\n" + f"morphism net : F -> T [{option}] ~ Normal\n" + "\n" + "program p : R -> R\n" + " observe y : R <- net(x)\n" + " return y\n" + "\n" + "export p\n" + ) + + +def test_param_source_name_selects_the_architecture() -> None: + """A bare name picks the source: ``linear`` is one matrix, not the + default MLP.""" + assert isinstance( + _family_of(_program("param_source=linear")).param_source, LinearSource + ) + assert isinstance(_family_of(_program("param_source=mlp")).param_source, MLPSource) + + +def test_param_source_call_form_carries_the_hidden_widths() -> None: + """``mlp(16, 8)`` is a call, not a bare name, and its arguments are + the hidden widths. The widths have to survive the option decode and + reach the source, so the built net is 16 then 8 rather than the + default 64, 64.""" + src = _family_of(_program("param_source=mlp(16, 8)")).param_source + assert isinstance(src, MLPSource) + widths = [layer.out_features for layer in src.net if hasattr(layer, "out_features")] + # Two hidden widths as written, then the family's param_dim (loc, log-scale). + assert widths == [16, 8, 2] + + default = _family_of(_program("param_source=mlp")).param_source + default_widths = [ + layer.out_features for layer in default.net if hasattr(layer, "out_features") + ] + assert default_widths == [64, 64, 2] + + +# --------------------------------------------------------------------------- +# Every family that has a parameter source can select it +# --------------------------------------------------------------------------- + + +_FAMILIES_WITHOUT_A_SOURCE = frozenset( + { + "ConditionalIndependent", + "ConditionalTransformed", + "ConditionalGaussianProcess", + "ConditionalHorseshoe", + } +) + + +def test_every_family_with_a_source_accepts_the_option() -> None: + """`[param_source=...]` is a morphism option the compiler accepts + for any family-backed kernel, so every family whose parameters come + from a source has to be able to receive it. A family that cannot is + a `TypeError` raised inside a constructor, with no line or column + on it.""" + import inspect + + from quivers.continuous import families + + missing = [] + for name, cls in vars(families).items(): + if not name.startswith("Conditional") or name in _FAMILIES_WITHOUT_A_SOURCE: + continue + try: + params = inspect.signature(cls.__init__).parameters + except TypeError, ValueError: # pragma: no cover - builtins + continue + if "param_source_option" not in params or "param_source" not in params: + missing.append(name) + assert not missing, f"families that cannot receive param_source: {missing}" + + +def _source_of(option: str, family: str = "Normal"): + from quivers.dsl import loads + + src = ( + "object F : Real 2\n" + "object T : Real 1\n" + "object R : FinSet 8\n" + "\n" + f"morphism net : F -> T {option} ~ {family}\n" + "\n" + "program p : R -> R\n" + " observe y : R <- net(x)\n" + " return y\n" + "\n" + "export p\n" + ) + model = loads(src).morphism + assert model is not None + return dict(model.named_modules())["_step_y._family"].param_source + + +def _widths(source) -> list[int]: + return [ + layer.out_features for layer in source.net if hasattr(layer, "out_features") + ] + + +@pytest.mark.parametrize("family", ["Normal", "Beta", "Gamma", "LogitNormal"]) +def test_param_source_reaches_every_family(family: str) -> None: + """Not just `Normal`: the option has to build the source whatever + the family.""" + assert isinstance(_source_of("[param_source=mlp]", family), MLPSource) + assert isinstance(_source_of("[param_source=linear]", family), LinearSource) + + +def test_hidden_dim_takes_a_sequence_of_widths() -> None: + """One entry per hidden layer, so the option says how many layers as + well as how wide. A single number cannot say the first.""" + source = _source_of("[param_source=mlp, hidden_dim=[64, 32, 16]]") + # three hidden layers as written, then the family's param_dim. + assert _widths(source) == [64, 32, 16, 2] + + +def test_hidden_dim_takes_a_single_width() -> None: + """A bare number is the one-layer case, so `n` reads as `[n]`.""" + assert _widths(_source_of("[param_source=mlp, hidden_dim=32]")) == [32, 2] + + +def test_hidden_dim_is_rejected_when_nothing_would_read_it() -> None: + """A single matrix has no hidden layers to be wide, so a width + against it is a question with no answer. Silence would read as one.""" + with pytest.raises(ValueError, match="no hidden layers"): + _source_of("[hidden_dim=32]") + + +def test_param_source_is_rejected_by_a_family_that_has_none() -> None: + """`Horseshoe` draws its own scale rather than reading parameters + from a source, so the option is refused at compile time with a + location, not inside a constructor.""" + from quivers.dsl import CompileError, loads + + src = ( + "object F : Real 2\n" + "object T : Real 2\n" + "object R : FinSet 8\n" + "\n" + "morphism net : F -> T [param_source=mlp] ~ Horseshoe\n" + "\n" + "program p : R -> R\n" + " observe y : R <- net(x)\n" + " return y\n" + "\n" + "export p\n" + ) + with pytest.raises(CompileError, match="does not take"): + loads(src) diff --git a/tests/test_plate.py b/tests/test_plate.py index 4ab5745d..44f20df4 100644 --- a/tests/test_plate.py +++ b/tests/test_plate.py @@ -18,6 +18,7 @@ from __future__ import annotations +import pytest import torch @@ -287,3 +288,99 @@ def test_parametric_template_morphism_param(self): """ c = self._compile(src) assert "demo" in c._morphisms + + +# --------------------------------------------------------------------------- +# Observed responses carry their codomain's event axis +# --------------------------------------------------------------------------- + + +_SCALAR_KERNEL = """object F : Real 1 +object T : Real 1 +object R : FinSet 5 + +morphism net : F -> T [param_source=linear] ~ Normal + +program p : R -> R + observe y : R <- net(x) + return y + +export p +""" + + +def test_scalar_response_scores_the_same_either_shape() -> None: + """A one-dimensional continuous codomain accepts N scalar responses + as ``(N,)`` or ``(N, 1)``. Both must score identically: the family's + parameters carry the event axis, so a bare ``(N,)`` response would + otherwise broadcast against ``(N, 1)`` into an ``(N, N)`` grid and + sum to a finite but wrong log-joint. + """ + from quivers.dsl import loads + + torch.manual_seed(0) + model = loads(_SCALAR_KERNEL).morphism + assert model is not None + x = torch.randn(5, 1) + flat = torch.randn(5) + column = flat.unsqueeze(-1) + + got_flat = model.log_joint(torch.zeros(5, 1), {"y": flat, "x": x})[0] + got_column = model.log_joint(torch.zeros(5, 1), {"y": column, "x": x})[0] + assert torch.allclose(got_flat, got_column, atol=1e-6) + + family = dict(model.named_modules())["_step_y._family"] + mu, sigma = family._get_params(x) + expected = torch.distributions.Normal(mu, sigma).log_prob(column).sum() + assert torch.allclose(got_flat, expected, atol=1e-5) + + +def test_response_whose_event_axis_is_wrong_is_rejected() -> None: + """A response that cannot carry the codomain's event shape is an + error, not something to broadcast into shape.""" + from quivers.continuous.families import ConditionalNormal + from quivers.continuous.plate import VectorisedObserve + from quivers.continuous.spaces import Euclidean + + domain = Euclidean(name="F", dim=1) + codomain = Euclidean(name="T", dim=3) + family = ConditionalNormal(domain, codomain) + obs = VectorisedObserve(family, torch.randn(5, 3)) + x = torch.randn(5, 1) + + # A 3-dimensional codomain has no reading of a bare (N,) response. + with pytest.raises(ValueError, match="expected \\(N, 3\\)"): + obs.log_prob(x, torch.randn(5)) + + # Nor of one whose event axis is the wrong width. + with pytest.raises(ValueError, match="event axis of size 3"): + obs.log_prob(x, torch.randn(5, 2)) + + +def test_inline_family_response_keeps_its_plate_shape() -> None: + """An inline family takes its parameters from the program's + environment as per-row values already aligned with the plate, so its + response indexes the plate and must not gain an event axis. + `OrderedLogistic` scores ordinal levels against a real-valued + predictor, and unsqueezing its index would push the score off by an + axis.""" + from quivers.dsl import loads + + torch.manual_seed(0) + program = loads( + "object Resp : FinSet 6\n" + "program ord : Resp -> Resp\n" + " observe y : Resp <- OrderedLogistic(eta, cuts)\n" + " return y\n" + "export ord\n" + ) + model = program.morphism + assert model is not None + n = 8 + obs = { + "y": torch.randint(0, 6, (n,)), + "eta": torch.randn(n), + "cuts": torch.linspace(-2.0, 2.0, 5), + } + out = model.log_joint(torch.zeros(n, 1), obs) + assert torch.isfinite(out).all() diff --git a/uv.lock b/uv.lock index 6fdc1649..83114a50 100644 --- a/uv.lock +++ b/uv.lock @@ -1421,7 +1421,7 @@ wheels = [ [[package]] name = "quivers" -version = "0.15.0" +version = "0.16.0" source = { editable = "." } dependencies = [ { name = "didactic" },