shardes 0.1.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (91) hide show
  1. shardes-0.1.1/.gitignore +8 -0
  2. shardes-0.1.1/CHANGELOG.md +33 -0
  3. shardes-0.1.1/CITATION.cff +18 -0
  4. shardes-0.1.1/CLAUDE.md +149 -0
  5. shardes-0.1.1/LICENSE +202 -0
  6. shardes-0.1.1/PKG-INFO +204 -0
  7. shardes-0.1.1/README.md +172 -0
  8. shardes-0.1.1/conftest.py +18 -0
  9. shardes-0.1.1/docs/00-context.md +219 -0
  10. shardes-0.1.1/docs/01-phase0-estimator-harness.md +579 -0
  11. shardes-0.1.1/docs/02-phase1-sharded-core.md +304 -0
  12. shardes-0.1.1/docs/04-phase3-coupling.md +256 -0
  13. shardes-0.1.1/docs/14-repo-split.md +309 -0
  14. shardes-0.1.1/docs/BACKLOG.md +217 -0
  15. shardes-0.1.1/docs/README.md +20 -0
  16. shardes-0.1.1/docs/compute.md +254 -0
  17. shardes-0.1.1/docs/conventions.md +276 -0
  18. shardes-0.1.1/docs/diagnosis-replicated-evaluation.md +279 -0
  19. shardes-0.1.1/docs/diagnosis-seed-regenerated-scan.md +164 -0
  20. shardes-0.1.1/docs/postmortem-gpu-determinism.md +781 -0
  21. shardes-0.1.1/docs/proposal-bf16-policy.md +173 -0
  22. shardes-0.1.1/docs/proposal-review-fixes.md +181 -0
  23. shardes-0.1.1/docs/proposal-scan-strategies-distribute.md +243 -0
  24. shardes-0.1.1/docs/provenance/README.md +20 -0
  25. shardes-0.1.1/docs/provenance/commit-map-2026-09-22.txt +491 -0
  26. shardes-0.1.1/examples/quickstart.py +34 -0
  27. shardes-0.1.1/pyproject.toml +85 -0
  28. shardes-0.1.1/src/shardes/__init__.py +56 -0
  29. shardes-0.1.1/src/shardes/check.py +172 -0
  30. shardes-0.1.1/src/shardes/contraction.py +150 -0
  31. shardes-0.1.1/src/shardes/core.py +443 -0
  32. shardes-0.1.1/src/shardes/coupling.py +324 -0
  33. shardes-0.1.1/src/shardes/dimensions.py +45 -0
  34. shardes-0.1.1/src/shardes/estimator.py +116 -0
  35. shardes-0.1.1/src/shardes/metrics.py +84 -0
  36. shardes-0.1.1/src/shardes/nn.py +108 -0
  37. shardes-0.1.1/src/shardes/problems/__init__.py +11 -0
  38. shardes-0.1.1/src/shardes/problems/control.py +172 -0
  39. shardes-0.1.1/src/shardes/problems/mlp.py +4 -0
  40. shardes-0.1.1/src/shardes/problems/quadratic.py +62 -0
  41. shardes-0.1.1/src/shardes/problems/qwen2.py +344 -0
  42. shardes-0.1.1/src/shardes/problems/transformer_block.py +111 -0
  43. shardes-0.1.1/src/shardes/py.typed +0 -0
  44. shardes-0.1.1/src/shardes/shaping.py +206 -0
  45. shardes-0.1.1/src/shardes/sharding.py +192 -0
  46. shardes-0.1.1/src/shardes/strategies/__init__.py +27 -0
  47. shardes-0.1.1/src/shardes/strategies/_noise.py +52 -0
  48. shardes-0.1.1/src/shardes/strategies/_scale.py +96 -0
  49. shardes-0.1.1/src/shardes/strategies/_select.py +50 -0
  50. shardes-0.1.1/src/shardes/strategies/iid_gaussian.py +118 -0
  51. shardes-0.1.1/src/shardes/strategies/lowrank.py +349 -0
  52. shardes-0.1.1/src/shardes/strategies/mirrored.py +153 -0
  53. shardes-0.1.1/src/shardes/strategies/protocol.py +118 -0
  54. shardes-0.1.1/src/shardes/strategies/registry.py +111 -0
  55. shardes-0.1.1/src/shardes/strategies/seed_regenerated.py +161 -0
  56. shardes-0.1.1/src/shardes/transforms/__init__.py +1 -0
  57. shardes-0.1.1/src/shardes/transforms/fwht.py +44 -0
  58. shardes-0.1.1/src/shardes/types.py +19 -0
  59. shardes-0.1.1/tests/conftest.py +71 -0
  60. shardes-0.1.1/tests/gpu/test_device_invariance_gpu.py +163 -0
  61. shardes-0.1.1/tests/test_accelerator_coverage.py +110 -0
  62. shardes-0.1.1/tests/test_bf16_policy.py +231 -0
  63. shardes-0.1.1/tests/test_check.py +153 -0
  64. shardes-0.1.1/tests/test_contraction.py +206 -0
  65. shardes-0.1.1/tests/test_control.py +157 -0
  66. shardes-0.1.1/tests/test_core.py +681 -0
  67. shardes-0.1.1/tests/test_coupling.py +544 -0
  68. shardes-0.1.1/tests/test_dimensions.py +67 -0
  69. shardes-0.1.1/tests/test_estimator.py +298 -0
  70. shardes-0.1.1/tests/test_examples.py +38 -0
  71. shardes-0.1.1/tests/test_fwht.py +98 -0
  72. shardes-0.1.1/tests/test_invariants.py +62 -0
  73. shardes-0.1.1/tests/test_lowrank.py +444 -0
  74. shardes-0.1.1/tests/test_metrics.py +149 -0
  75. shardes-0.1.1/tests/test_mirrored.py +144 -0
  76. shardes-0.1.1/tests/test_nn.py +174 -0
  77. shardes-0.1.1/tests/test_no_dangling_paths.py +47 -0
  78. shardes-0.1.1/tests/test_problems.py +85 -0
  79. shardes-0.1.1/tests/test_public_api.py +126 -0
  80. shardes-0.1.1/tests/test_qwen2.py +145 -0
  81. shardes-0.1.1/tests/test_qwen2_golden.py +122 -0
  82. shardes-0.1.1/tests/test_registry.py +133 -0
  83. shardes-0.1.1/tests/test_shaping.py +297 -0
  84. shardes-0.1.1/tests/test_sharding.py +409 -0
  85. shardes-0.1.1/tests/test_strategies.py +415 -0
  86. shardes-0.1.1/tests/test_transformer_block.py +103 -0
  87. shardes-0.1.1/validation/README.md +18 -0
  88. shardes-0.1.1/validation/kaggle/t2prime/kernel-metadata.json +14 -0
  89. shardes-0.1.1/validation/kaggle/t2prime/t2prime.py +71 -0
  90. shardes-0.1.1/validation/reference.json +259 -0
  91. shardes-0.1.1/validation/reference.py +145 -0
@@ -0,0 +1,8 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ build/
5
+ dist/
6
+ .pytest_cache/
7
+ .ipynb_checkpoints/
8
+ .venv/
@@ -0,0 +1,33 @@
1
+ # Changelog
2
+
3
+ The supported surface is `shardes.__all__`. Before 1.0 a minor version may change it, and
4
+ says so here. Everything else is importable and may move without notice.
5
+
6
+ ## 0.1.1
7
+
8
+ Packaging only; the code is 0.1.0's.
9
+
10
+ - First release on PyPI: `pip install shardes`.
11
+ - Index metadata (classifiers, keywords, links), and an explicit list of what the sdist
12
+ contains.
13
+ - A release workflow: it builds from the tag, checks that the tag names the declared
14
+ version, checks the artifacts, installs the wheel into a clean environment, and only
15
+ then publishes, through PyPI's trusted publishing rather than a token.
16
+
17
+ ## 0.1.0
18
+
19
+ The first version with a repository of its own. Until the tag `monorepo-final` the library
20
+ shared one with the paper it was built for, which is now
21
+ [shardes-paper](https://github.com/andreshernandez-spec/shardes-paper).
22
+
23
+ - `ShardedES`: `init`, `ask`, `apply`, `tell`, over the one-axis device mesh `make_mesh`
24
+ builds.
25
+ - Strategies: `IIDGaussian`, `SeedRegenerated(chunk)`, `LowRank(r)`, and `Mirrored(inner)`
26
+ for antithetic pairs around any of them.
27
+ - Two placements of the update contraction, `how="A"` and `how="B"`.
28
+ - Shaping: `centered_ranks` (the default), `centered`, `group_relative`, `none`.
29
+ - The seams a model routes its weights through, `shardes.nn.dense` and `shardes.nn.embed`,
30
+ and `shardes.check.check_model` to find the weights that miss them.
31
+ - Problems: `quadratic`, `mlp`, `transformer_block`, and a Qwen2.5 port that needs
32
+ `shardes[models]`.
33
+ - Python 3.12 or newer, JAX 0.11 or newer.
@@ -0,0 +1,18 @@
1
+ # How to cite the library. GitHub reads this file and shows a "Cite this repository"
2
+ # button. `version` and `date-released` name the latest release, not main, so they move
3
+ # when a tag is cut, and a test holds `version` to the changelog's latest released heading.
4
+ cff-version: 1.2.0
5
+ message: "If you use shardes, please cite it."
6
+ type: software
7
+ title: "shardes: sharded evolution strategies for JAX"
8
+ authors:
9
+ - family-names: Hernandez
10
+ given-names: Andres
11
+ version: 0.1.1
12
+ date-released: "2026-09-22"
13
+ repository-code: "https://github.com/andreshernandez-spec/shardes"
14
+ license: Apache-2.0
15
+ keywords:
16
+ - evolution strategies
17
+ - JAX
18
+ - distributed training
@@ -0,0 +1,149 @@
1
+ # CLAUDE.md
2
+
3
+ Instructions for Claude Code working in this repository. Keep this file short — it
4
+ loads into every session. Detail lives in `docs/`.
5
+
6
+ ---
7
+
8
+ ## What this project is
9
+
10
+ `shardes` is a JAX library for **sharded evolution strategies**: an `ask`/`eval`/`tell`
11
+ core that shards the population and the rollouts across devices, with the perturbation
12
+ scheme as a pluggable, shape-aware strategy rather than a hardcoded `randn(num_dims)`.
13
+
14
+ The distribution state is replicated, not sharded (decided in `docs/02` C1.4 and corrected
15
+ here on 2026-07-31). Sharding it is theatre for isotropic ES and, for the CMA family, would
16
+ cost a per-generation gather to save memory the design already spends on replicating the
17
+ model.
18
+
19
+ It exists because the two 2025 papers that made ES work at LLM scale
20
+ ([Qiu et al.](https://arxiv.org/abs/2509.24372), full-rank + seed regeneration;
21
+ [Sarkar et al. / EGGROLL](https://arxiv.org/abs/2511.16652), rank-`r` factored) have no
22
+ common library, and the incumbent (`evosax`) forecloses both by flattening every solution
23
+ to one dense vector via `ravel_pytree`. Background: `docs/00-context.md`.
24
+
25
+ **The deliverable is the library.** Everything else is instrumentation for it.
26
+
27
+ ---
28
+
29
+ ## Ground rules
30
+
31
+ ### 1. Andres writes and understands the load-bearing code
32
+
33
+ This is a hard constraint, not a preference. The project's value is what he can explain
34
+ in a technical interview, and the ecosystem this targets is explicit about it — JAX's
35
+ `docs/contributing.md` states that *"the use of an AI agent that autonomously writes code
36
+ and submits pull requests is not permitted"* and *"do not use AI to speak for you"* on
37
+ issue trackers. This repo is standalone, so the letter of that policy doesn't bind it,
38
+ but the spirit is the whole point.
39
+
40
+ Concretely:
41
+
42
+ - **Claude Code may**: scaffold, write tests, write benchmark harnesses, write plotting
43
+ and infra code, port reference implementations for comparison, research, review, explain,
44
+ find bugs, refactor.
45
+ - **Claude Code should hand back**: the perturbation strategies, the sharding logic, the
46
+ gradient-estimator math, and anything that ends up in a README or a talk. Draft it,
47
+ explain the reasoning, then let Andres write the version that ships.
48
+ - **Claude Code must never**: draft GitHub issues, PR descriptions, or review comments for
49
+ upstream projects. Andres writes those in his own voice.
50
+
51
+ If asked to do something in the second or third bucket, say so and offer the first-bucket
52
+ version instead.
53
+
54
+ ### 2. Every number in a doc has a script behind it
55
+
56
+ No claim of the form "X is Y× faster" or "error is 1e-15" enters any markdown file without
57
+ a committed, re-runnable script and a recorded environment (GPU model, JAX version, commit
58
+ SHA). Scripts that check the library are in `tools/`; the experiments and the paper are in
59
+ the [shardes-paper](https://github.com/andreshernandez-spec/shardes-paper) repository, and a
60
+ number that comes from there is cited as `shardes-paper:<path>`. If a number can't be
61
+ reproduced from a clean checkout, delete it.
62
+
63
+ ### 3. Don't build what a measurement already ruled out
64
+
65
+ Phase 3, coupled sampling at scale, was conditional on a Phase 0 result and that result
66
+ came back no (`docs/04-phase3-coupling.md`, `docs/01`). Do not build it speculatively
67
+ because it seems interesting. The research program's phases and gates are in
68
+ `shardes-paper:PLAN.md`.
69
+
70
+ ---
71
+
72
+ ## Repository layout
73
+
74
+ ```
75
+ shardes/
76
+ ├── README.md for a user: install, quickstart, guarantees
77
+ ├── CLAUDE.md this file
78
+ ├── pyproject.toml must stay pip-installable from a git SHA; shardes-paper pins one
79
+ ├── src/shardes/ the library
80
+ ├── tests/ pytest; CPU-only, no network; tests/gpu needs real accelerators
81
+ ├── examples/ what the README shows; the suite runs them
82
+ ├── tools/ scripts that check the library, not experiments
83
+ ├── validation/ the real-hardware invariance check (gate G1)
84
+ └── docs/ design (02), context (00), the estimator study (01), conventions,
85
+ diagnoses, proposals, and the record of the split (14)
86
+ ```
87
+
88
+ The experiments, the results, the paper and the campaign docs are in
89
+ [shardes-paper](https://github.com/andreshernandez-spec/shardes-paper), which installs this
90
+ library at a pinned commit. Until the tag `monorepo-final` they were here, and that history
91
+ is kept on purpose: records from that period cite commits of this repository. Never rewrite
92
+ it, and never delete a `provenance/*` tag: each one keeps a cited commit reachable. It was
93
+ rewritten once, on 2026-09-22, to remove a personal account from three files; every commit
94
+ after 2026-08-01 changed hash and `docs/provenance/` maps old to new. That is the exception,
95
+ and the map is what it cost.
96
+
97
+ `tests/` runs on CPU, no GPU, no network, in **two tiers**: `pytest --fast` is the inner
98
+ loop while editing, `pytest` is everything and is the default. Budgets and the reasoning
99
+ are in `docs/conventions.md`. This line said "under two minutes" until 2026-08-01; that
100
+ number predated the code and was about to cost `R` in the unbiasedness tests.
101
+
102
+ ---
103
+
104
+ ## Environment
105
+
106
+ - Python ≥ 3.12 (jax 0.11 requires it; this said 3.11 until CI tried it), JAX ≥ 0.11. **Do not pin below 0.11** — `from jax import shard_map` needs
107
+ 0.8 and `AxisType` needs 0.11, and those are what the library is built on. This used to
108
+ say "evosax is stuck at `<0.7`"; that stopped being true at evosax 0.2.0 (`jax>=0.5.0`,
109
+ no upper bound). The floor is justified by what we use, not by what they pin.
110
+ - `from jax import shard_map`. `jax.experimental.shard_map` is deprecated (JAX 0.8.0).
111
+ - Multi-device logic is developed and tested **on CPU** with
112
+ `XLA_FLAGS=--xla_force_host_platform_device_count=8`. See `docs/compute.md`. Do not
113
+ reach for a GPU to debug a `PartitionSpec`.
114
+ - Dependencies stay minimal: `jax`, `numpy`, `pytest`. Adding anything else needs a
115
+ one-line justification in the PR description. `optax`, `flax`, `chex` are acceptable if
116
+ actually used; `evosax` is a comparison target, not a dependency.
117
+
118
+ ---
119
+
120
+ ## Invariants — breaking these is a bug, not a tradeoff
121
+
122
+ 1. **No global flattening.** Nothing in `src/` calls `ravel_pytree` on a solution. Leaves
123
+ keep their `(m, n)` shape all the way through sample → apply → contract. This is the
124
+ architectural difference from evosax and the reason low-rank perturbation is expressible
125
+ at all.
126
+ 2. **Device-count invariance.** For a fixed seed, the update produced on 1 device and on
127
+ 8 devices must agree to within float tolerance (`rtol=1e-5` bf16, `1e-12` f32). Seeds
128
+ are derived from the *member index*, never from the device index. There is a test for
129
+ this; it is the most important test in the repo.
130
+ 3. **The perturbation is never materialized by the low-rank path.** If a profile shows an
131
+ `(n_members, m, n)` array being allocated under `LowRank`, the implementation is wrong.
132
+ 4. **Communication is measured, not assumed.** Every `psum`/`all_gather` in the update path
133
+ is accounted for in `shardes-paper:docs/03-phase2-benchmarks.md`. See the note there about the two
134
+ contraction strategies — the "ES only all-reduces scalars" claim is true only for one
135
+ of them.
136
+
137
+ ---
138
+
139
+ ## Working style
140
+
141
+ - Small commits, each with its test. Prefer a failing test first.
142
+ - When a design decision has two defensible answers, write both down in the relevant
143
+ `docs/` file with the tradeoff, and flag it for Andres rather than picking silently.
144
+ - Numerical code: assert against an exact oracle where one exists (see
145
+ `docs/conventions.md` for the list — e.g. the FWHT has one that ships in JAX).
146
+ - The names in `shardes.__all__` are the supported surface and `tests/test_public_api.py`
147
+ pins them. Changing one is a version bump and a line in the release notes. Everything
148
+ else stays importable from its module, and shardes-paper uses those deep paths, so a
149
+ rename there breaks the paper's pin when it next moves.
shardes-0.1.1/LICENSE ADDED
@@ -0,0 +1,202 @@
1
+
2
+ Apache License
3
+ Version 2.0, January 2004
4
+ http://www.apache.org/licenses/
5
+
6
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
7
+
8
+ 1. Definitions.
9
+
10
+ "License" shall mean the terms and conditions for use, reproduction,
11
+ and distribution as defined by Sections 1 through 9 of this document.
12
+
13
+ "Licensor" shall mean the copyright owner or entity authorized by
14
+ the copyright owner that is granting the License.
15
+
16
+ "Legal Entity" shall mean the union of the acting entity and all
17
+ other entities that control, are controlled by, or are under common
18
+ control with that entity. For the purposes of this definition,
19
+ "control" means (i) the power, direct or indirect, to cause the
20
+ direction or management of such entity, whether by contract or
21
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
22
+ outstanding shares, or (iii) beneficial ownership of such entity.
23
+
24
+ "You" (or "Your") shall mean an individual or Legal Entity
25
+ exercising permissions granted by this License.
26
+
27
+ "Source" form shall mean the preferred form for making modifications,
28
+ including but not limited to software source code, documentation
29
+ source, and configuration files.
30
+
31
+ "Object" form shall mean any form resulting from mechanical
32
+ transformation or translation of a Source form, including but
33
+ not limited to compiled object code, generated documentation,
34
+ and conversions to other media types.
35
+
36
+ "Work" shall mean the work of authorship, whether in Source or
37
+ Object form, made available under the License, as indicated by a
38
+ copyright notice that is included in or attached to the work
39
+ (an example is provided in the Appendix below).
40
+
41
+ "Derivative Works" shall mean any work, whether in Source or Object
42
+ form, that is based on (or derived from) the Work and for which the
43
+ editorial revisions, annotations, elaborations, or other modifications
44
+ represent, as a whole, an original work of authorship. For the purposes
45
+ of this License, Derivative Works shall not include works that remain
46
+ separable from, or merely link (or bind by name) to the interfaces of,
47
+ the Work and Derivative Works thereof.
48
+
49
+ "Contribution" shall mean any work of authorship, including
50
+ the original version of the Work and any modifications or additions
51
+ to that Work or Derivative Works thereof, that is intentionally
52
+ submitted to Licensor for inclusion in the Work by the copyright owner
53
+ or by an individual or Legal Entity authorized to submit on behalf of
54
+ the copyright owner. For the purposes of this definition, "submitted"
55
+ means any form of electronic, verbal, or written communication sent
56
+ to the Licensor or its representatives, including but not limited to
57
+ communication on electronic mailing lists, source code control systems,
58
+ and issue tracking systems that are managed by, or on behalf of, the
59
+ Licensor for the purpose of discussing and improving the Work, but
60
+ excluding communication that is conspicuously marked or otherwise
61
+ designated in writing by the copyright owner as "Not a Contribution."
62
+
63
+ "Contributor" shall mean Licensor and any individual or Legal Entity
64
+ on behalf of whom a Contribution has been received by Licensor and
65
+ subsequently incorporated within the Work.
66
+
67
+ 2. Grant of Copyright License. Subject to the terms and conditions of
68
+ this License, each Contributor hereby grants to You a perpetual,
69
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
70
+ copyright license to reproduce, prepare Derivative Works of,
71
+ publicly display, publicly perform, sublicense, and distribute the
72
+ Work and such Derivative Works in Source or Object form.
73
+
74
+ 3. Grant of Patent License. Subject to the terms and conditions of
75
+ this License, each Contributor hereby grants to You a perpetual,
76
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
77
+ (except as stated in this section) patent license to make, have made,
78
+ use, offer to sell, sell, import, and otherwise transfer the Work,
79
+ where such license applies only to those patent claims licensable
80
+ by such Contributor that are necessarily infringed by their
81
+ Contribution(s) alone or by combination of their Contribution(s)
82
+ with the Work to which such Contribution(s) was submitted. If You
83
+ institute patent litigation against any entity (including a
84
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
85
+ or a Contribution incorporated within the Work constitutes direct
86
+ or contributory patent infringement, then any patent licenses
87
+ granted to You under this License for that Work shall terminate
88
+ as of the date such litigation is filed.
89
+
90
+ 4. Redistribution. You may reproduce and distribute copies of the
91
+ Work or Derivative Works thereof in any medium, with or without
92
+ modifications, and in Source or Object form, provided that You
93
+ meet the following conditions:
94
+
95
+ (a) You must give any other recipients of the Work or
96
+ Derivative Works a copy of this License; and
97
+
98
+ (b) You must cause any modified files to carry prominent notices
99
+ stating that You changed the files; and
100
+
101
+ (c) You must retain, in the Source form of any Derivative Works
102
+ that You distribute, all copyright, patent, trademark, and
103
+ attribution notices from the Source form of the Work,
104
+ excluding those notices that do not pertain to any part of
105
+ the Derivative Works; and
106
+
107
+ (d) If the Work includes a "NOTICE" text file as part of its
108
+ distribution, then any Derivative Works that You distribute must
109
+ include a readable copy of the attribution notices contained
110
+ within such NOTICE file, excluding those notices that do not
111
+ pertain to any part of the Derivative Works, in at least one
112
+ of the following places: within a NOTICE text file distributed
113
+ as part of the Derivative Works; within the Source form or
114
+ documentation, if provided along with the Derivative Works; or,
115
+ within a display generated by the Derivative Works, if and
116
+ wherever such third-party notices normally appear. The contents
117
+ of the NOTICE file are for informational purposes only and
118
+ do not modify the License. You may add Your own attribution
119
+ notices within Derivative Works that You distribute, alongside
120
+ or as an addendum to the NOTICE text from the Work, provided
121
+ that such additional attribution notices cannot be construed
122
+ as modifying the License.
123
+
124
+ You may add Your own copyright statement to Your modifications and
125
+ may provide additional or different license terms and conditions
126
+ for use, reproduction, or distribution of Your modifications, or
127
+ for any such Derivative Works as a whole, provided Your use,
128
+ reproduction, and distribution of the Work otherwise complies with
129
+ the conditions stated in this License.
130
+
131
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
132
+ any Contribution intentionally submitted for inclusion in the Work
133
+ by You to the Licensor shall be under the terms and conditions of
134
+ this License, without any additional terms or conditions.
135
+ Notwithstanding the above, nothing herein shall supersede or modify
136
+ the terms of any separate license agreement you may have executed
137
+ with Licensor regarding such Contributions.
138
+
139
+ 6. Trademarks. This License does not grant permission to use the trade
140
+ names, trademarks, service marks, or product names of the Licensor,
141
+ except as required for reasonable and customary use in describing the
142
+ origin of the Work and reproducing the content of the NOTICE file.
143
+
144
+ 7. Disclaimer of Warranty. Unless required by applicable law or
145
+ agreed to in writing, Licensor provides the Work (and each
146
+ Contributor provides its Contributions) on an "AS IS" BASIS,
147
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
148
+ implied, including, without limitation, any warranties or conditions
149
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
150
+ PARTICULAR PURPOSE. You are solely responsible for determining the
151
+ appropriateness of using or redistributing the Work and assume any
152
+ risks associated with Your exercise of permissions under this License.
153
+
154
+ 8. Limitation of Liability. In no event and under no legal theory,
155
+ whether in tort (including negligence), contract, or otherwise,
156
+ unless required by applicable law (such as deliberate and grossly
157
+ negligent acts) or agreed to in writing, shall any Contributor be
158
+ liable to You for damages, including any direct, indirect, special,
159
+ incidental, or consequential damages of any character arising as a
160
+ result of this License or out of the use or inability to use the
161
+ Work (including but not limited to damages for loss of goodwill,
162
+ work stoppage, computer failure or malfunction, or any and all
163
+ other commercial damages or losses), even if such Contributor
164
+ has been advised of the possibility of such damages.
165
+
166
+ 9. Accepting Warranty or Additional Liability. While redistributing
167
+ the Work or Derivative Works thereof, You may choose to offer,
168
+ and charge a fee for, acceptance of support, warranty, indemnity,
169
+ or other liability obligations and/or rights consistent with this
170
+ License. However, in accepting such obligations, You may act only
171
+ on Your own behalf and on Your sole responsibility, not on behalf
172
+ of any other Contributor, and only if You agree to indemnify,
173
+ defend, and hold each Contributor harmless for any liability
174
+ incurred by, or claims asserted against, such Contributor by reason
175
+ of your accepting any such warranty or additional liability.
176
+
177
+ END OF TERMS AND CONDITIONS
178
+
179
+ APPENDIX: How to apply the Apache License to your work.
180
+
181
+ To apply the Apache License to your work, attach the following
182
+ boilerplate notice, with the fields enclosed by brackets "[]"
183
+ replaced with your own identifying information. (Don't include
184
+ the brackets!) The text should be enclosed in the appropriate
185
+ comment syntax for the file format. We also recommend that a
186
+ file or class name and description of purpose be included on the
187
+ same "printed page" as the copyright notice for easier
188
+ identification within third-party archives.
189
+
190
+ Copyright 2026 Andres Hernandez
191
+
192
+ Licensed under the Apache License, Version 2.0 (the "License");
193
+ you may not use this file except in compliance with the License.
194
+ You may obtain a copy of the License at
195
+
196
+ http://www.apache.org/licenses/LICENSE-2.0
197
+
198
+ Unless required by applicable law or agreed to in writing, software
199
+ distributed under the License is distributed on an "AS IS" BASIS,
200
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
201
+ See the License for the specific language governing permissions and
202
+ limitations under the License.
shardes-0.1.1/PKG-INFO ADDED
@@ -0,0 +1,204 @@
1
+ Metadata-Version: 2.5
2
+ Name: shardes
3
+ Version: 0.1.1
4
+ Summary: Sharded evolution strategies for JAX
5
+ Project-URL: Repository, https://github.com/andreshernandez-spec/shardes
6
+ Project-URL: Changelog, https://github.com/andreshernandez-spec/shardes/blob/main/CHANGELOG.md
7
+ Project-URL: Issues, https://github.com/andreshernandez-spec/shardes/issues
8
+ Project-URL: Paper, https://github.com/andreshernandez-spec/shardes-paper
9
+ Author: Andres Hernandez
10
+ License-Expression: Apache-2.0
11
+ License-File: LICENSE
12
+ Keywords: distributed,evolution strategies,jax,low-rank,sharding
13
+ Classifier: Development Status :: 3 - Alpha
14
+ Classifier: Intended Audience :: Science/Research
15
+ Classifier: Operating System :: OS Independent
16
+ Classifier: Programming Language :: Python :: 3 :: Only
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
20
+ Classifier: Typing :: Typed
21
+ Requires-Python: >=3.12
22
+ Requires-Dist: jax>=0.11
23
+ Requires-Dist: numpy
24
+ Requires-Dist: scipy
25
+ Provides-Extra: dev
26
+ Requires-Dist: pytest; extra == 'dev'
27
+ Provides-Extra: models
28
+ Requires-Dist: safetensors; extra == 'models'
29
+ Provides-Extra: tasks
30
+ Requires-Dist: playground; extra == 'tasks'
31
+ Description-Content-Type: text/markdown
32
+
33
+ # shardes
34
+
35
+ Sharded evolution strategies for JAX.
36
+
37
+ Evolution strategies train a model without gradients: perturb the weights, score each
38
+ perturbed copy, move toward the copies that scored well. `shardes` is an `ask` / `apply` /
39
+ `tell` core that splits that population across devices, with the perturbation scheme as a
40
+ pluggable strategy. The two 2025 results that made ES work at language-model scale,
41
+ full-rank noise regenerated from seeds ([Qiu et al.](https://arxiv.org/abs/2509.24372))
42
+ and rank-`r` factored noise ([EGGROLL](https://arxiv.org/abs/2511.16652)), are one
43
+ constructor argument apart here, on the same mesh, shaping and update.
44
+
45
+ > **Status: 0.1, pre-alpha.** The API can still change. What is here is tested, on CPU with
46
+ > eight simulated devices, and device-count invariance has been checked on real hardware
47
+ > too: the GPU suite on two T4s, and one device against eight A100s on a 0.5B-parameter
48
+ > model.
49
+
50
+ ## Install
51
+
52
+ ```sh
53
+ pip install shardes # the release
54
+ pip install "shardes @ git+https://github.com/andreshernandez-spec/shardes" # main
55
+ ```
56
+
57
+ Python 3.12 or newer, JAX 0.11 or newer. On an accelerator, install JAX for it first
58
+ (`pip install -U "jax[cuda12]"` or `"jax[tpu]"`). The core needs only jax, numpy and scipy.
59
+ Extras: `shardes[models]` for loading Qwen2.5 checkpoints, `shardes[tasks]` for the MuJoCo
60
+ Playground adapter.
61
+
62
+ ## Quickstart
63
+
64
+ This is `examples/quickstart.py`, and the test suite runs it. It needs no accelerator:
65
+ eight simulated CPU devices stand in, and the sharding is the same program either way.
66
+
67
+ ```python
68
+ import os
69
+ os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8" # 8 devices on a CPU
70
+
71
+ import jax
72
+ import shardes
73
+ from shardes.problems import transformer_block # a small model that uses the seams
74
+
75
+ key = jax.random.key(0)
76
+ params = transformer_block.init(key, d_model=64)
77
+ batch = transformer_block.make_batch(jax.random.fold_in(key, 1), d_model=64, batch=8, seq=32)
78
+
79
+ mesh = shardes.make_mesh() # every visible device, one "pop" axis
80
+ es = shardes.ShardedES(shardes.Mirrored(shardes.LowRank(r=1)),
81
+ n=256, sigma=1e-2, lr=1e-2, mesh=mesh)
82
+ state = es.init(key, params)
83
+
84
+ @jax.jit
85
+ def generation(state):
86
+ pert, state = es.ask(state) # a Perturbation, never a batch of params
87
+ fitness = es.apply(transformer_block.loss, state, pert)(batch) # one loss per member
88
+ return es.tell(state, pert, fitness), fitness
89
+
90
+ for step in range(30):
91
+ state, fitness = generation(state)
92
+ if step % 10 == 0 or step == 29:
93
+ print(f"step {step:2d} mean loss {float(fitness.mean()):.4f}")
94
+ ```
95
+
96
+ Jit the whole generation rather than stepping it eagerly; that is what lets JAX settle
97
+ device placement at trace time. `tell` descends, so hand it a loss, or the negative of a
98
+ reward.
99
+
100
+ ## Both published algorithms, one argument apart
101
+
102
+ ```python
103
+ # Qiu et al. 2025: full-rank noise, regenerated from seeds, small population
104
+ shardes.ShardedES(shardes.Mirrored(shardes.SeedRegenerated()), n=30, ...)
105
+
106
+ # EGGROLL (Sarkar et al. 2025): rank-1 factors, never materialized, huge population
107
+ shardes.ShardedES(shardes.Mirrored(shardes.LowRank(r=1)), n=262_144, ...)
108
+ ```
109
+
110
+ | strategy | the perturbation | what it costs |
111
+ |---|---|---|
112
+ | `IIDGaussian()` | full rank, materialized | memory: `n` copies of the noise |
113
+ | `SeedRegenerated(chunk=1)` | full rank, regenerated from each member's seed | compute: every perturbation is drawn twice. `chunk` trades memory for speed |
114
+ | `LowRank(r)` | rank `r`, two thin factors per matrix, product never formed | the noise is not full rank |
115
+ | `Mirrored(inner)` | antithetic pairs around any of the above | halves the distinct directions |
116
+
117
+ `ask` returns a `Perturbation`, not a batch of parameter trees, and the rest follows from
118
+ that: under `LowRank` it is a pair of factors, under `SeedRegenerated` a key and member
119
+ ids. Shaping is `centered_ranks` by default (`shardes.shaping` has `centered`,
120
+ `group_relative` and `none`).
121
+
122
+ ## Where the update is assembled: `how="A"` or `how="B"`
123
+
124
+ After evaluation each device holds the fitnesses of its own members, and the update
125
+ `sum_i w_i eps_i` has to be put together. There are two placements, and which is faster
126
+ depends on the perturbation and on the interconnect.
127
+
128
+ - **`how="B"`, the default.** Each device contracts its own members, then the partial
129
+ updates are all-reduced. That moves a buffer the size of the model.
130
+ - **`how="A"`.** Every device gathers the scalar fitnesses and contracts the whole
131
+ population itself. Only scalars cross the wire, and the contraction is repeated on
132
+ every device.
133
+
134
+ Measured on an 8x A100 node, a TPU v5e-8 and two A100 nodes over sockets: on one host
135
+ with a fast interconnect, **B wins for full-rank and seed-regenerated perturbations, and A
136
+ wins for low-rank perturbations on large models**, by more on the TPU. Across a slow host
137
+ boundary A wins more widely. The measurements, and the paper they belong to, are in
138
+ [shardes-paper](https://github.com/andreshernandez-spec/shardes-paper).
139
+
140
+ ## Your own model
141
+
142
+ A low-rank perturbation is never materialized, so a perturbed weight is not an array: it
143
+ is a base matrix plus two factors. A model therefore routes its parameterized matrix
144
+ multiplies and embedding lookups through two seams, `shardes.nn.dense(x, w)` and
145
+ `shardes.nn.embed(table, ids)`, which do the right thing for a plain array and for a
146
+ structured weight alike. Direct arithmetic on a structured weight raises rather than
147
+ silently densifying. `shardes.problems.transformer_block` is the small worked example and
148
+ `shardes.problems.qwen2` is Qwen2.5 ported this way.
149
+ `shardes.check.check_model(model, params, batch)` finds, in a second on CPU, every weight a
150
+ model reaches without a seam, instead of the strategy raising minutes into a run.
151
+
152
+ ## What it guarantees
153
+
154
+ - **Device count cannot change the result.** Member `i`'s noise derives from `i` alone,
155
+ so the update contracted on one device and on eight agrees to floating-point
156
+ tolerance. This is tested on simulated devices and checked on real ones: `tests/gpu`
157
+ on two T4s (`validation/`), and on A100s the update for Qwen2.5-0.5B computed on one
158
+ device and on eight agrees to 6.3e-6 relative error
159
+ (`shardes-paper:experiments/countdown/results/c6d-a100x8-2026-08-18`).
160
+ - **The low-rank path never forms an `(n, m, n)` array.** A test inspects the traced
161
+ program to make sure.
162
+ - **Fitness is float32 or wider.** In bfloat16 a population's losses collapse to a
163
+ handful of ties and rank shaping of ties is noise, so `tell` refuses rather than casts.
164
+ - **No global flattening.** Parameters keep their shapes from sampling to update, which
165
+ is what makes per-matrix structure expressible at all.
166
+
167
+ ## Development
168
+
169
+ ```sh
170
+ pip install -e ".[dev]"
171
+ pytest --fast # the inner loop
172
+ pytest # everything, including the statistical tier
173
+ ```
174
+
175
+ The suite pins JAX to the CPU and simulates eight devices, so it needs no accelerator and
176
+ no network. `tools/` holds scripts that check the library itself (mutation testing,
177
+ compile-cost diagnostics, two probes behind design decisions), `validation/` the
178
+ real-hardware invariance check, and `docs/` the design, conventions and diagnoses.
179
+ CI runs the suite on Python 3.12 and 3.13 and installs the built wheel into a clean
180
+ environment.
181
+
182
+ ## The paper and the experiments
183
+
184
+ This library is the instrument of *Update-contraction placement in sharded evolution
185
+ strategies on GPUs and TPUs*. The manuscript, every experiment and every result live in
186
+ [shardes-paper](https://github.com/andreshernandez-spec/shardes-paper), which installs
187
+ this library at a pinned commit.
188
+
189
+ Until the tag `monorepo-final` the two were one repository, and that history is kept here
190
+ on purpose: result records from that period cite commits of this repository, and
191
+ checking one out gives the driver, its config and the library together, as run.
192
+
193
+ ## Non-goals
194
+
195
+ - Sharded *parameters*. Every device holds the model and evaluates independently;
196
+ sharding parameters would bring back the communication ES avoids.
197
+ - A general replacement for evosax. This targets the sharded, large-population,
198
+ structured-perturbation regime.
199
+ - Reimplementing either paper's full experimental setup. The papers stand; this is
200
+ infrastructure.
201
+
202
+ ## License
203
+
204
+ Apache 2.0.