anyinit 0.1.0__tar.gz

anyinit-0.1.0/.gitignore

@@ -0,0 +1,14 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+build/
+dist/
+*.egg-info/
+.venv/
+data/
+_old_code/
+site/

anyinit-0.1.0/CHANGELOG.md

@@ -0,0 +1,30 @@
+# Changelog
+
+## 0.1.0
+
+First release.
+
+### Added
+
+- `initialize(model, mode, input_spec, ...)` — one call initializes any model, with
+  the framework detected from it. The remaining options are explicit keywords:
+  `distribution`, `center`, `gains`, `seed` and `params`. PyTorch, Keras 3 (on any of
+  its backends) and JAX (via Flax), each an optional dependency, so installing one never
+  imports another.
+- `initialize_params(model, params, ...)` for functional frameworks, which return a new
+  parameter tree rather than mutating one.
+- Two solve modes. `analytic` propagates moments through the graph and never runs the
+  model; `empirical` measures real batches and assumes nothing distributional. They share
+  the graph, the layer adapters and the target policy.
+- `register_activation` for arbitrary activations. Their moment maps are computed by
+  Gauss–Legendre quadrature over the Gaussian, exact to machine precision across kinks.
+- `anyinit.gain(name, **params)`: the gain AnyInit gives an activation, `√2` for ReLU.
+  `initialize(gains={...})` fixes the gain for an activation instead of solving for it.
+- Depth-stability diagnosis. Every activation gets `χ = d log E[f²] / d log Var[z]`, which
+  equals the homogeneity degree for homogeneous activations and so identifies the ones no
+  scalar initialization can stabilize.
+- `InitReport`: the scale chosen for every layer, the graph fidelity actually achieved, a
+  measured validation pass sized against its own sampling noise, and anything AnyInit
+  could not do. `assert_healthy()` raises on the findings.
+- Weights are drawn in NumPy in a canonical layout from the run's seed, so the same
+  architecture and seed produce bit-identical weights under all three frameworks.

anyinit-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jose I. Mestre
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

anyinit-0.1.0/PKG-INFO

@@ -0,0 +1,161 @@
+Metadata-Version: 2.5
+Name: anyinit
+Version: 0.1.0
+Summary: Initialize any model, in any framework, correctly — with one call
+Project-URL: Homepage, https://github.com/jmiravet/AnyInit
+Project-URL: Documentation, https://jmiravet.github.io/AnyInit/
+Project-URL: Repository, https://github.com/jmiravet/AnyInit
+Project-URL: Issues, https://github.com/jmiravet/AnyInit/issues
+Project-URL: Changelog, https://github.com/jmiravet/AnyInit/blob/main/CHANGELOG.md
+Author-email: "Jose I. Mestre" <jmiravet@uji.es>, Alberto Fernández-Hernández <a.fernandez@upv.es>
+Maintainer-email: "Jose I. Mestre" <jmiravet@uji.es>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: deep-learning,flax,initialization,jax,keras,pytorch,signal-propagation,tensorflow,weight-initialization
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.21
+Provides-Extra: dev
+Requires-Dist: mypy; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pytest>=7; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# AnyInit
+
+[![CI](https://github.com/jmiravet/AnyInit/actions/workflows/ci.yml/badge.svg)](https://github.com/jmiravet/AnyInit/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/jmiravet/AnyInit/graph/badge.svg)](https://codecov.io/gh/jmiravet/AnyInit)
+[![PyPI](https://img.shields.io/pypi/v/anyinit)](https://pypi.org/project/anyinit/)
+[![Python](https://img.shields.io/pypi/pyversions/anyinit)](https://pypi.org/project/anyinit/)
+[![Docs](https://img.shields.io/badge/docs-jmiravet.github.io%2FAnyInit-blue)](https://jmiravet.github.io/AnyInit/)
+[![License: MIT](https://img.shields.io/badge/license-MIT-blue)](https://github.com/jmiravet/AnyInit/blob/main/LICENSE)
+
+Initialize any model, in any framework, correctly — with one call.
+
+```python
+import anyinit
+
+report = anyinit.initialize(model)
+print(report)
+```
+
+`model` can be a PyTorch module, a Keras model or a Flax module. AnyInit traces it, finds
+which activation follows each layer, and scales every weight so the signal neither dies
+nor explodes on its way through. Nothing to configure, nothing to look up.
+
+## Install
+
+```bash
+pip install anyinit
+```
+
+AnyInit depends on NumPy alone and uses whichever framework your model comes from.
+
+## What "any" means
+
+**Any framework.** PyTorch, Keras 3 (on TensorFlow, JAX or PyTorch) and Flax, detected
+from the model. The same seed gives the same weights in all three.
+
+**Any activation.** AnyInit measures the function itself instead of looking it up in a
+table, so a new activation is one decorator away:
+
+```python
+@anyinit.register_activation
+class ReLUCubed(torch.nn.Module):
+    def forward(self, x):
+        return torch.relu(x) ** 3
+```
+
+**Any architecture.** Branches, residual additions, concatenations, normalization,
+pooling, dropout, attention and embeddings.
+
+**Theory or measurement.** Scale from theory, with no data, or from real batches:
+
+```python
+anyinit.initialize(model, "analytic", (32, 3, 224, 224))   # shapes only, milliseconds
+anyinit.initialize(model, "empirical", batch)              # measures, assumes nothing
+```
+
+## It tells you when it cannot help
+
+Some activations cannot be kept stable across depth by any initialization. AnyInit says
+so in the report instead of handing back a network that will not train:
+
+```
+Stability
+  relu3   chi= 3.000  sigma*=  0.7148  unstable degree=3
+      ! relu3 is homogeneous of degree 3 (chi=3.000), so a relative error grows by
+        3.00x per layer and reaches 3.49e+09x over 20 layers. No scalar initialization
+        is depth-stable here: reduce depth, insert normalization, or use a degree-one
+        activation
+```
+
+Call `report.assert_healthy()` to turn that into an exception, for example in CI.
+
+## Learn more
+
+The [documentation](https://jmiravet.github.io/AnyInit/) covers:
+
+- [Reading the report](https://jmiravet.github.io/AnyInit/guide/report/)
+- [Analytic or empirical](https://jmiravet.github.io/AnyInit/guide/modes/)
+- [Stability across depth](https://jmiravet.github.io/AnyInit/guide/stability/)
+- [Frameworks](https://jmiravet.github.io/AnyInit/guide/frameworks/)
+- [Options](https://jmiravet.github.io/AnyInit/guide/options/)
+- [Why a gain table is not enough](https://jmiravet.github.io/AnyInit/why/), with measurements
+- [Tutorials](https://jmiravet.github.io/AnyInit/tutorials/custom-activation/) and the
+  [API](https://jmiravet.github.io/AnyInit/api/)
+
+Contributions are welcome; see
+[Contributing](https://jmiravet.github.io/AnyInit/contributing/).
+
+## References
+
+Variance scaling comes from Glorot and Bengio (2010) and He et al. (2015), whose rectifier
+gain falls out of the moment map as the degree-one case. Treating signal propagation as a
+dynamical system, and the order parameter behind the `χ` diagnostic, come from Poole et al.
+(2016) and Schoenholz et al. (2017). The empirical mode generalizes LSUV (Mishkin & Matas,
+2016) from a sequence of layers to a graph. Fixed points of the variance map, and the SELU
+constants the reference table reproduces, are from Klambauer et al. (2017). The
+`sinusoidal` distribution is from Fernández-Hernández et al. (2025).
+
+What is new is evaluating the moment map numerically, by Gaussian quadrature, instead of
+deriving it per activation, which is why AnyInit accepts an activation it has never seen.
+
+Fernández-Hernández, A., Mestre, J. I., Dolz, M. F., Duato, J., & Quintana-Ortí, E. S.
+(2025). Sinusoidal initialization, time for a new start. In *Advances in Neural
+Information Processing Systems* (Vol. 38). https://doi.org/10.48550/arXiv.2505.12909
+
+Glorot, X., & Bengio, Y. (2010). Understanding the difficulty of training deep feedforward
+neural networks. In *Proceedings of the Thirteenth International Conference on Artificial
+Intelligence and Statistics* (pp. 249–256). PMLR.
+https://proceedings.mlr.press/v9/glorot10a.html
+
+He, K., Zhang, X., Ren, S., & Sun, J. (2015). Delving deep into rectifiers: Surpassing
+human-level performance on ImageNet classification. In *Proceedings of the IEEE
+International Conference on Computer Vision* (pp. 1026–1034). IEEE.
+https://doi.org/10.1109/ICCV.2015.123
+
+Klambauer, G., Unterthiner, T., Mayr, A., & Hochreiter, S. (2017). Self-normalizing neural
+networks. In *Advances in Neural Information Processing Systems* (Vol. 30, pp. 971–980).
+https://arxiv.org/abs/1706.02515
+
+Mishkin, D., & Matas, J. (2016). All you need is a good init. In *International Conference
+on Learning Representations*. https://arxiv.org/abs/1511.06422
+
+Poole, B., Lahiri, S., Raghu, M., Sohl-Dickstein, J., & Ganguli, S. (2016). Exponential
+expressivity in deep neural networks through transient chaos. In *Advances in Neural
+Information Processing Systems* (Vol. 29, pp. 3360–3368).
+https://arxiv.org/abs/1606.05340
+
+Schoenholz, S. S., Gilmer, J., Ganguli, S., & Sohl-Dickstein, J. (2017). Deep information
+propagation. In *International Conference on Learning Representations*.
+https://arxiv.org/abs/1611.01232

anyinit-0.1.0/README.md

@@ -0,0 +1,129 @@
+# AnyInit
+
+[![CI](https://github.com/jmiravet/AnyInit/actions/workflows/ci.yml/badge.svg)](https://github.com/jmiravet/AnyInit/actions/workflows/ci.yml)
+[![codecov](https://codecov.io/gh/jmiravet/AnyInit/graph/badge.svg)](https://codecov.io/gh/jmiravet/AnyInit)
+[![PyPI](https://img.shields.io/pypi/v/anyinit)](https://pypi.org/project/anyinit/)
+[![Python](https://img.shields.io/pypi/pyversions/anyinit)](https://pypi.org/project/anyinit/)
+[![Docs](https://img.shields.io/badge/docs-jmiravet.github.io%2FAnyInit-blue)](https://jmiravet.github.io/AnyInit/)
+[![License: MIT](https://img.shields.io/badge/license-MIT-blue)](https://github.com/jmiravet/AnyInit/blob/main/LICENSE)
+
+Initialize any model, in any framework, correctly — with one call.
+
+```python
+import anyinit
+
+report = anyinit.initialize(model)
+print(report)
+```
+
+`model` can be a PyTorch module, a Keras model or a Flax module. AnyInit traces it, finds
+which activation follows each layer, and scales every weight so the signal neither dies
+nor explodes on its way through. Nothing to configure, nothing to look up.
+
+## Install
+
+```bash
+pip install anyinit
+```
+
+AnyInit depends on NumPy alone and uses whichever framework your model comes from.
+
+## What "any" means
+
+**Any framework.** PyTorch, Keras 3 (on TensorFlow, JAX or PyTorch) and Flax, detected
+from the model. The same seed gives the same weights in all three.
+
+**Any activation.** AnyInit measures the function itself instead of looking it up in a
+table, so a new activation is one decorator away:
+
+```python
+@anyinit.register_activation
+class ReLUCubed(torch.nn.Module):
+    def forward(self, x):
+        return torch.relu(x) ** 3
+```
+
+**Any architecture.** Branches, residual additions, concatenations, normalization,
+pooling, dropout, attention and embeddings.
+
+**Theory or measurement.** Scale from theory, with no data, or from real batches:
+
+```python
+anyinit.initialize(model, "analytic", (32, 3, 224, 224))   # shapes only, milliseconds
+anyinit.initialize(model, "empirical", batch)              # measures, assumes nothing
+```
+
+## It tells you when it cannot help
+
+Some activations cannot be kept stable across depth by any initialization. AnyInit says
+so in the report instead of handing back a network that will not train:
+
+```
+Stability
+  relu3   chi= 3.000  sigma*=  0.7148  unstable degree=3
+      ! relu3 is homogeneous of degree 3 (chi=3.000), so a relative error grows by
+        3.00x per layer and reaches 3.49e+09x over 20 layers. No scalar initialization
+        is depth-stable here: reduce depth, insert normalization, or use a degree-one
+        activation
+```
+
+Call `report.assert_healthy()` to turn that into an exception, for example in CI.
+
+## Learn more
+
+The [documentation](https://jmiravet.github.io/AnyInit/) covers:
+
+- [Reading the report](https://jmiravet.github.io/AnyInit/guide/report/)
+- [Analytic or empirical](https://jmiravet.github.io/AnyInit/guide/modes/)
+- [Stability across depth](https://jmiravet.github.io/AnyInit/guide/stability/)
+- [Frameworks](https://jmiravet.github.io/AnyInit/guide/frameworks/)
+- [Options](https://jmiravet.github.io/AnyInit/guide/options/)
+- [Why a gain table is not enough](https://jmiravet.github.io/AnyInit/why/), with measurements
+- [Tutorials](https://jmiravet.github.io/AnyInit/tutorials/custom-activation/) and the
+  [API](https://jmiravet.github.io/AnyInit/api/)
+
+Contributions are welcome; see
+[Contributing](https://jmiravet.github.io/AnyInit/contributing/).
+
+## References
+
+Variance scaling comes from Glorot and Bengio (2010) and He et al. (2015), whose rectifier
+gain falls out of the moment map as the degree-one case. Treating signal propagation as a
+dynamical system, and the order parameter behind the `χ` diagnostic, come from Poole et al.
+(2016) and Schoenholz et al. (2017). The empirical mode generalizes LSUV (Mishkin & Matas,
+2016) from a sequence of layers to a graph. Fixed points of the variance map, and the SELU
+constants the reference table reproduces, are from Klambauer et al. (2017). The
+`sinusoidal` distribution is from Fernández-Hernández et al. (2025).
+
+What is new is evaluating the moment map numerically, by Gaussian quadrature, instead of
+deriving it per activation, which is why AnyInit accepts an activation it has never seen.
+
+Fernández-Hernández, A., Mestre, J. I., Dolz, M. F., Duato, J., & Quintana-Ortí, E. S.
+(2025). Sinusoidal initialization, time for a new start. In *Advances in Neural
+Information Processing Systems* (Vol. 38). https://doi.org/10.48550/arXiv.2505.12909
+
+Glorot, X., & Bengio, Y. (2010). Understanding the difficulty of training deep feedforward
+neural networks. In *Proceedings of the Thirteenth International Conference on Artificial
+Intelligence and Statistics* (pp. 249–256). PMLR.
+https://proceedings.mlr.press/v9/glorot10a.html
+
+He, K., Zhang, X., Ren, S., & Sun, J. (2015). Delving deep into rectifiers: Surpassing
+human-level performance on ImageNet classification. In *Proceedings of the IEEE
+International Conference on Computer Vision* (pp. 1026–1034). IEEE.
+https://doi.org/10.1109/ICCV.2015.123
+
+Klambauer, G., Unterthiner, T., Mayr, A., & Hochreiter, S. (2017). Self-normalizing neural
+networks. In *Advances in Neural Information Processing Systems* (Vol. 30, pp. 971–980).
+https://arxiv.org/abs/1706.02515
+
+Mishkin, D., & Matas, J. (2016). All you need is a good init. In *International Conference
+on Learning Representations*. https://arxiv.org/abs/1511.06422
+
+Poole, B., Lahiri, S., Raghu, M., Sohl-Dickstein, J., & Ganguli, S. (2016). Exponential
+expressivity in deep neural networks through transient chaos. In *Advances in Neural
+Information Processing Systems* (Vol. 29, pp. 3360–3368).
+https://arxiv.org/abs/1606.05340
+
+Schoenholz, S. S., Gilmer, J., Ganguli, S., & Sohl-Dickstein, J. (2017). Deep information
+propagation. In *International Conference on Learning Representations*.
+https://arxiv.org/abs/1611.01232

anyinit-0.1.0/docs/api.md

@@ -0,0 +1,44 @@
+# API
+
+Everything below is importable from the top-level `anyinit` package.
+
+## Initializing
+
+::: anyinit.initialize
+
+::: anyinit.initialize_params
+
+## Activations
+
+::: anyinit.register_activation
+
+::: anyinit.unregister_activation
+
+::: anyinit.registered_activations
+
+::: anyinit.gain
+
+::: anyinit.activation_profile
+
+::: anyinit.ActivationProfile
+
+## The report
+
+::: anyinit.InitReport
+
+::: anyinit.LayerRecord
+
+::: anyinit.StabilityRecord
+
+## Backends
+
+::: anyinit.available_backends
+
+::: anyinit.known_backends
+
+## Errors
+
+::: anyinit.errors
+    options:
+      show_root_heading: false
+      members_order: source

anyinit-0.1.0/docs/changelog.md

@@ -0,0 +1 @@
+--8<-- "CHANGELOG.md"

anyinit-0.1.0/docs/contributing.md

@@ -0,0 +1,40 @@
+# Contributing
+
+## Setup
+
+```bash
+pip install -e '.[dev]' torch tensorflow jax flax
+pytest                                  # backend tests skip when a framework is absent
+ruff check src tests examples docs
+ruff format --check src tests examples docs
+mypy src                                # strict
+
+pip install -r docs/requirements.txt
+mkdocs serve                            # the documentation site, locally
+```
+
+[Design](design.md) explains how the code is laid out.
+
+## Adding a backend
+
+Implement `anyinit.backends.base.Backend` and add it to `_BACKENDS` in
+`anyinit/backends/__init__.py`. None of the numerics is touched. The contract is:
+
+- `handles(model)` — recognize the model from its class hierarchy, without importing the
+  framework. `module_roots` does the inspection.
+- `build_graph(model, input_spec)` — trace into a `ModelGraph`.
+- `eval_elementwise(fn, x)` — apply a native callable to quadrature abscissas.
+- `read_weight` / `write_weight` / `write_bias` / `write_gain` — canonical layout in and
+  out. `write_gain` assigns rather than multiplies, so `initialize` stays idempotent.
+- `forward_taps(model, inputs, taps, training=True)` — one forward pass, returning moments
+  at the named nodes. `training` must normally be true: at initialization a batch-norm
+  layer's running statistics are still (0, 1), so in inference mode it scales by gamma
+  instead of normalizing. Running statistics must be restored afterward, and random state
+  fixed during the pass, so that dropout draws the same masks every time and the empirical
+  solver sees a deterministic function of the scales.
+- `unscaled_weights(model, graph)` — every weight no graph node writes, with the reason.
+- `label` — how the report names the backend; Keras adds what it runs on, `keras[jax]`.
+- `finalize(model)` — new parameters for a functional framework, `None` otherwise.
+
+`tests/test_cross_backend.py` then checks that the new backend produces the same weights
+as the others from the same seed.

anyinit-0.1.0/docs/design.md

@@ -0,0 +1,104 @@
+# Design
+
+Notes for anyone reading or extending the code. To add a framework, see
+[Contributing](contributing.md#adding-a-backend).
+
+## Two layers
+
+```
+┌──────────────────────────────────────────────────────┐
+│ anyinit.core — imports no framework                 │
+│ quadrature · profiles · stability · graph IR ·       │
+│ topology · transfer · fan · distributions · solvers  │
+└───────────────────────┬──────────────────────────────┘
+                        │  ModelGraph + numpy.ndarray
+┌───────────────────────┴──────────────────────────────┐
+│ anyinit.backends — one module per framework         │
+│ pytorch.py · keras.py · jax.py                       │
+│ tracing · weight layouts · forward passes · taps     │
+└──────────────────────────────────────────────────────┘
+```
+
+The core never holds a tensor. It speaks `ModelGraph` — an IR of typed nodes — and
+`numpy.ndarray` in a canonical layout. Backends trace the native model into that IR,
+translate layouts, read and write parameters, and run forward passes returning
+*statistics* rather than tensors.
+
+## Principles
+
+1. **The core imports no framework.** `tests/test_core_isolation.py` walks the AST of
+   every core module to enforce it, and runs a subprocess to check that importing
+   `anyinit` loads none of them.
+2. **Lazy imports.** A framework is imported inside the method that needs it, never at
+   module scope, so installing one never pulls in another.
+3. **Reinitialize, not rescale.** The analytic recursion assumes `W` is i.i.d. zero-mean
+   with variance `σ_w²`; drawing the weights makes that true by construction.
+4. **One graph, two solvers.** `analytic` and `empirical` share the IR, the topology, the
+   layer adapters, the target policy and the report. They differ only in where the
+   moments come from.
+5. **Mode purity.** No measurement ever sets an `analytic` scale. The model is run only
+   for the validation pass, which runs when `input_spec` is given and whose result
+   reaches the report alone, and, under Flax, for a 16-row probe that identifies
+   activations (below). Shapes come from fake tensors,
+   which carry no data. `empirical` never consults a Gaussian profile to set a scale.
+6. **Weights drawn in NumPy.** The same seed gives the same weights under every backend.
+7. **Fail visibly.** Anything AnyInit cannot do correctly goes into the report rather
+   than being approximated silently.
+
+## Canonical weight layout
+
+The core reasons in `(fan_out, fan_in, *receptive_field)`. Each backend translates:
+
+| Layer | Canonical | PyTorch | Keras | Flax |
+|---|---|---|---|---|
+| Dense | `(out, in)` | `(out, in)` | `(in, out)` | `(in, out)` |
+| Conv2D | `(out, in/g, kh, kw)` | `(out, in/g, kh, kw)` | `(kh, kw, in, out)` | `(kh, kw, in, out)` |
+| ConvTranspose2D | `(out, in/g, kh, kw)` | `(in, out/g, kh, kw)` | `(kh, kw, out, in)` | `(kh, kw, in, out)` |
+| Embedding | `(num, dim)` | `(num, dim)` | `(num, dim)` | `(num, dim)` |
+
+A transposed convolution stores its input and output axes the other way round from an
+ordinary one, so fans computed on the native shape come out inverted.
+
+## Tracing
+
+| Framework | Mechanism | Fidelity |
+|---|---|---|
+| PyTorch | `torch.fx` symbolic trace; shapes from fake-tensor propagation | graph; falls back to a linear chain from `named_modules` when tracing fails |
+| Keras | `model.operations` and each operation's inbound nodes | graph; linear for an unwired Sequential |
+| Flax | Flax layers and `jax.nn` activations wrapped inside a context manager, then the model applied to a 16-row probe | linear |
+
+PyTorch's `nn.Transformer*` layers are traced as single modules, so the adapter expands
+each into its real structure — attention, residual additions, normalizations and the
+feed-forward block — with both `norm_first` layouts. The expanded nodes have no FX node of
+their own and are measured through hooks on the submodules that produce or consume them.
+Adaptive pooling records no window; the fake-tensor shapes supply it.
+
+A jaxpr is too far from the source to read activations off: most `jax.nn` functions lower
+to opaque wrappers and a bias add is indistinguishable from a residual one. The
+instrumented trace recovers call order instead, and each layer reports `scope.path`, which
+is its key in the parameter tree. That yields an ordered chain rather than a DAG, since
+nothing intercepts `+`, and the graph is marked `linear`. An activation that never passes
+through `jax.nn` — a module field, a closure, a lambda — is not seen being called, so the
+probe's values identify it: wherever one layer's output reaches the next transformed, the
+transform is matched against the known and registered activations.
+
+Whatever a backend reaches but cannot scale — recurrent weights, attention layers it has no
+adapter for, bare parameters — is listed in the report with the reason, never skipped
+silently.
+
+## Where things live
+
+| Concern | Module |
+|---|---|
+| Gaussian integration, kink detection | `core/quadrature.py` |
+| NumPy reference activations | `core/activations.py` |
+| Moment map of an activation | `core/profile.py` |
+| Activation registry | `core/registry.py` |
+| `χ`, feasibility, verdicts | `core/stability.py` |
+| Typed graph IR | `core/graph.py` |
+| Layer/activation pairing | `core/topology.py` |
+| Moment transfer per node kind | `core/transfer.py` |
+| Canonical shapes and fan arithmetic | `core/fan.py` |
+| Weight samplers | `core/distributions.py` |
+| Solvers | `core/analytic.py`, `core/empirical.py` |
+| Orchestration | `_run.py` |

anyinit-0.1.0/docs/experiments/README.md

@@ -0,0 +1,21 @@
+# Reproducing the measurements
+
+The figures quoted in the guide, in `../why.md` and in `../reference.md` come from these.
+They are kept out of the test suite because they take minutes rather than seconds.
+
+| script | what it measures |
+|---|---|
+| `quadrature_accuracy.py` | Gauss-Legendre against Gauss-Hermite and Monte Carlo on `relu³`, whose moments are known exactly |
+| `gaussian_breakdown.py` | How far a high moment departs from the Gaussian prediction with width, and whether the Gamma correction tracks it |
+| `stability_table.py` | `χ` and `σ*` for every builtin activation and for `relu^p` |
+| `gain_table.py` | `E[a²]` through a 20-layer MLP under the usual gain table and under both AnyInit modes |
+| `draw_drift.py` | How far one draw of a deep ReLU MLP lands from the analytic prediction, with and without `center=True`, against the empirical mode |
+| `depth_sweep.py` | Measured rescaling across depth and width for `relu^p`, showing which degrees survive |
+
+Run any of them directly:
+
+```bash
+python docs/experiments/stability_table.py
+```
+
+`depth_sweep.py`, `draw_drift.py` and `gain_table.py` need PyTorch; the others need only NumPy.

anyinit-0.1.0/docs/experiments/depth_sweep.py

@@ -0,0 +1,80 @@
+"""Which activations survive depth, measured rather than predicted.
+
+Rescales each layer from measurements so that E[a^2] is one after every activation -- the
+best any per-layer scalar can do -- and reports what actually arrives at the end.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+
+DEPTHS_AND_WIDTHS = ((3, 256), (5, 256), (10, 256), (20, 256), (20, 1024), (20, 4096))
+
+
+def network(power: int, depth: int, width: int) -> tuple[nn.Module, type[nn.Module]]:
+    class Power(nn.Module):
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return torch.relu(x) ** power
+
+    layers: list[nn.Module] = []
+    for _ in range(depth):
+        layers += [nn.Linear(width, width, bias=False), Power()]
+    return nn.Sequential(*layers), Power
+
+
+def rescale(model: nn.Module, activation: type[nn.Module], power: int, rows: int) -> None:
+    """Drive E[a^2] to one after every activation, measuring as we go."""
+    first = next(m for m in model if isinstance(m, nn.Linear))
+    hidden = torch.randn(rows, first.in_features)
+    with torch.no_grad():
+        for index, layer in enumerate(model):
+            if isinstance(layer, nn.Linear):
+                follower = model[index + 1]
+                for _ in range(40):
+                    value = float(follower(layer(hidden)).pow(2).mean())
+                    if not math.isfinite(value) or value <= 0.0:
+                        layer.weight.mul_(0.5)
+                        continue
+                    layer.weight.mul_((1.0 / value) ** (1.0 / (2 * power)))
+                    if abs(value - 1.0) < 1e-5:
+                        break
+            hidden = layer(hidden)
+
+
+def main(rows: int = 2048) -> None:
+    print(f"{'p':>3s} {'depth':>6s} {'width':>6s} {'final E[a^2]':>13s} {'top-8 energy':>13s}")
+    for power in (1, 2, 3):
+        for depth, width in DEPTHS_AND_WIDTHS:
+            torch.manual_seed(0)
+            model, activation = network(power, depth, width)
+            for module in model.modules():
+                if isinstance(module, nn.Linear):
+                    nn.init.normal_(module.weight, 0.0, 0.05)
+            rescale(model, activation, power, rows)
+
+            with torch.no_grad():
+                output = model(torch.randn(rows, width))
+            energy = output.pow(2)
+            total = float(energy.sum())
+            share = (
+                float(energy.sum(0).sort(descending=True).values[:8].sum()) / total
+                if total > 0.0
+                else float("nan")
+            )
+            print(
+                f"{power:3d} {depth:6d} {width:6d} {float(energy.mean()):13.4g} "
+                f"{share * 100:12.1f}%"
+            )
+
+    print(
+        "\nDegree one is stable at any depth and widening only spreads the energy further.\n"
+        "Above it no amount of per-layer rescaling holds: the signal concentrates into a\n"
+        "handful of units and then dies, and extra width does not help."
+    )
+
+
+if __name__ == "__main__":
+    main()