weightspace 0.2.0__tar.gz → 0.3.0__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.
- weightspace-0.3.0/CHANGELOG.md +57 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/CITATION.cff +1 -1
- {weightspace-0.2.0 → weightspace-0.3.0}/PKG-INFO +56 -2
- {weightspace-0.2.0 → weightspace-0.3.0}/README.md +52 -1
- weightspace-0.3.0/benchmarks/profile_hotpaths.py +111 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/docs/README.md +1 -0
- weightspace-0.3.0/docs/production.md +118 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/llms.txt +3 -1
- {weightspace-0.2.0 → weightspace-0.3.0}/pyproject.toml +3 -2
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/__init__.py +12 -2
- weightspace-0.3.0/src/weightspace/_chunks.py +101 -0
- weightspace-0.3.0/src/weightspace/agent.py +562 -0
- weightspace-0.3.0/src/weightspace/cli.py +125 -0
- weightspace-0.3.0/src/weightspace/events.py +153 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/llm/anthropic.py +23 -3
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/llm/base.py +5 -0
- weightspace-0.3.0/src/weightspace/llm/mock.py +61 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/llm/openai.py +49 -3
- weightspace-0.3.0/src/weightspace/otel.py +133 -0
- weightspace-0.3.0/src/weightspace/py.typed +0 -0
- weightspace-0.3.0/src/weightspace/render.py +255 -0
- weightspace-0.3.0/src/weightspace/result.py +103 -0
- weightspace-0.3.0/src/weightspace/sandbox.py +130 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/space.py +99 -39
- weightspace-0.3.0/src/weightspace/tools.py +36 -0
- weightspace-0.3.0/src/weightspace/usage.py +87 -0
- weightspace-0.3.0/tests/test_perf_paths.py +84 -0
- weightspace-0.3.0/tests/test_production.py +259 -0
- weightspace-0.2.0/CHANGELOG.md +0 -31
- weightspace-0.2.0/src/weightspace/agent.py +0 -446
- weightspace-0.2.0/src/weightspace/cli.py +0 -76
- weightspace-0.2.0/src/weightspace/llm/mock.py +0 -38
- weightspace-0.2.0/src/weightspace/render.py +0 -153
- weightspace-0.2.0/src/weightspace/sandbox.py +0 -70
- {weightspace-0.2.0 → weightspace-0.3.0}/.github/ISSUE_TEMPLATE/bug_report.yml +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.github/ISSUE_TEMPLATE/config.yml +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.github/ISSUE_TEMPLATE/feature_request.yml +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.github/dependabot.yml +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.github/pull_request_template.md +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.github/workflows/ci.yml +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.github/workflows/publish.yml +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.gitignore +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/.pre-commit-config.yaml +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/CODE_OF_CONDUCT.md +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/CONTRIBUTING.md +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/LICENSE +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/NOTICE +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/SECURITY.md +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/docs/concepts.md +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/docs/probes.md +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/docs/providers.md +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/examples/fix_quadrants.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/examples/hf_reduce_loops.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/examples/read_which_task.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/adapters.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/llm/__init__.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/probes.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/prompts.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/src/weightspace/tasks.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/tests/conftest.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/tests/test_agent.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/tests/test_llm.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/tests/test_render.py +0 -0
- {weightspace-0.2.0 → weightspace-0.3.0}/tests/test_space.py +0 -0
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
# Changelog
|
|
2
|
+
|
|
3
|
+
All notable changes to this project are documented here. The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project uses [Semantic Versioning](https://semver.org/).
|
|
4
|
+
|
|
5
|
+
## [Unreleased]
|
|
6
|
+
|
|
7
|
+
## [0.3.0] - 2026-11-03
|
|
8
|
+
|
|
9
|
+
### Added
|
|
10
|
+
- Run events and callbacks (`callbacks=`): `LoggingCallback`, `EventLog`, `TranscriptWriter`, `ConsoleLogger`, and `weightspace.otel.OpenTelemetryCallback` with GenAI semantic-convention spans (`weightspace[otel]`).
|
|
11
|
+
- `Usage` (tokens normalized across providers, incl. cache reads/writes and reasoning) and `Pricing`; `Result.tokens`, `cost`, `stop_reason`, `run_id`, `seconds`, `timings`, `error`, `ok`.
|
|
12
|
+
- `token_budget`, `cost_budget` and `time_budget`: the LLM gets a last turn to submit when one runs out.
|
|
13
|
+
- Ctrl-C returns a partial `Result`; API failures raise `SessionError` carrying the partial result.
|
|
14
|
+
- Read-only tool calls of one turn run in parallel (`parallel_tools=`, `wsp.tool(..., parallel=True)`).
|
|
15
|
+
- `ScriptedLLM.from_transcript` to replay logged runs offline; CLI `--replay`.
|
|
16
|
+
- OpenAI `prompt_cache_key` per run on api.openai.com; `max_retries` / `timeout` on all backends; Anthropic streaming for long requests.
|
|
17
|
+
- CLI: `--probe FILE.py:FN`, `--forward-budget`, `--price`, `--max-cost`, `--max-tokens`, `--max-time`, `--events`, `--json`, `--version`; exit code 2 when a run does not finish cleanly.
|
|
18
|
+
- `py.typed` marker.
|
|
19
|
+
- Exp. 8 results in the README: base, the paper's checkpoint and the Opus and Astra edits, with thinking on and off.
|
|
20
|
+
- `experiments/exp8`: `evaluate.py --no-thinking` and `table.py`.
|
|
21
|
+
|
|
22
|
+
### Changed
|
|
23
|
+
- `ws.diff()` and `ws.patch()` run in blocks (about 6x faster and 50x less memory on large tensors); `describe_tensor` is cached and uses a Gram-matrix spectrum for tall/wide matrices.
|
|
24
|
+
- The sandbox timeout works off the main thread and without SIGALRM; sandbox code runs under `torch.no_grad()`.
|
|
25
|
+
- `agent.py` split into `tools.py`, `result.py`, `events.py` and `usage.py` (same public API).
|
|
26
|
+
- Faster hot paths, measured with the new `benchmarks/profile_hotpaths.py`: block-parallel `diff`, `patch`, snapshots and restores; exact histogram-based quantiles; sparse checkpoints; content-keyed `describe` cache; cheaper spectrum for large matrices. `describe` on a 151936 x 1024 embedding goes from 6.5 s and 2.4 GB to 1.2 s and 132 MB, and a 30-step agent loop from 9.2 s to 0.4 s. See docs/production.md.
|
|
27
|
+
|
|
28
|
+
### Fixed
|
|
29
|
+
- `apply_patch` silently did nothing on tensors stored as non-contiguous arrays.
|
|
30
|
+
- `diff` and `patch` counted untouched NaN weights as edits.
|
|
31
|
+
|
|
32
|
+
## [0.2.0] - 2026-11-03
|
|
33
|
+
|
|
34
|
+
### Added
|
|
35
|
+
- `openresponses:<model>` backend (`OpenResponses`) for the [Open Responses](https://www.openresponses.org) spec (OpenAI Responses API), replaying encrypted reasoning between turns. `openai-responses:` is kept as an alias.
|
|
36
|
+
- `speed=` on `Session`, `wsp.read/fix/write/improve` and backends (CLI `--speed` / `--fast`) for provider fast mode: Anthropic `speed`, OpenAI `service_tier`. Saved in `result.json`.
|
|
37
|
+
- `usage.thinking_tokens` is recorded for Anthropic models.
|
|
38
|
+
- `llms.txt` with a short usage guide for coding agents.
|
|
39
|
+
- Exp. 8 replication harness for Qwen3.5-0.8B in `experiments/exp8`.
|
|
40
|
+
- Open-source project files, docs and a wider CI matrix.
|
|
41
|
+
|
|
42
|
+
### Changed
|
|
43
|
+
- Anthropic `thinking=False` falls back to `effort="low"` on models that can't disable thinking (Opus 5.5).
|
|
44
|
+
|
|
45
|
+
## [0.1.0] - 2026-11-03
|
|
46
|
+
|
|
47
|
+
### Added
|
|
48
|
+
- `WeightSpace` with NumPy, PyTorch, Hugging Face, SafeTensors and `.pt` stores; tracked edits, diffs, checkpoints and sparse patches.
|
|
49
|
+
- `read`, `fix`, `write` and `improve` sessions with a persistent Python sandbox, budgeted probes and text mode.
|
|
50
|
+
- OpenAI-compatible (OpenAI, OpenRouter, vLLM, SGLang, Ollama) and Anthropic backends.
|
|
51
|
+
- `thinking` switch (off / low / medium / high / xhigh / max), with adaptive thinking and prompt caching on Anthropic.
|
|
52
|
+
- `weightspace` CLI and examples.
|
|
53
|
+
|
|
54
|
+
[Unreleased]: https://github.com/leofmarciano/weightspace/compare/v0.3.0...HEAD
|
|
55
|
+
[0.3.0]: https://github.com/leofmarciano/weightspace/compare/v0.2.0...v0.3.0
|
|
56
|
+
[0.2.0]: https://github.com/leofmarciano/weightspace/compare/v0.1.0...v0.2.0
|
|
57
|
+
[0.1.0]: https://github.com/leofmarciano/weightspace/releases/tag/v0.1.0
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: weightspace
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: Weight-space programming: let an LLM read and write neural network weights directly.
|
|
5
5
|
Project-URL: Homepage, https://marcianotheory.com
|
|
6
6
|
Project-URL: Repository, https://github.com/leofmarciano/weightspace
|
|
@@ -27,6 +27,7 @@ Provides-Extra: anthropic
|
|
|
27
27
|
Requires-Dist: anthropic>=0.40; extra == 'anthropic'
|
|
28
28
|
Provides-Extra: dev
|
|
29
29
|
Requires-Dist: mypy>=1.10; extra == 'dev'
|
|
30
|
+
Requires-Dist: opentelemetry-sdk>=1.20; extra == 'dev'
|
|
30
31
|
Requires-Dist: pre-commit>=3; extra == 'dev'
|
|
31
32
|
Requires-Dist: pytest>=8; extra == 'dev'
|
|
32
33
|
Requires-Dist: ruff>=0.5; extra == 'dev'
|
|
@@ -36,6 +37,8 @@ Requires-Dist: torch>=2.1; extra == 'hf'
|
|
|
36
37
|
Requires-Dist: transformers>=4.40; extra == 'hf'
|
|
37
38
|
Provides-Extra: openai
|
|
38
39
|
Requires-Dist: openai>=1.40; extra == 'openai'
|
|
40
|
+
Provides-Extra: otel
|
|
41
|
+
Requires-Dist: opentelemetry-api>=1.20; extra == 'otel'
|
|
39
42
|
Provides-Extra: torch
|
|
40
43
|
Requires-Dist: torch>=2.1; extra == 'torch'
|
|
41
44
|
Description-Content-Type: text/markdown
|
|
@@ -182,7 +185,23 @@ res = s.run("...")
|
|
|
182
185
|
```
|
|
183
186
|
|
|
184
187
|
Custom tools: `wsp.tool("name", "description", x={"type": "number"})(fn)`.
|
|
185
|
-
Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script
|
|
188
|
+
Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script, and
|
|
189
|
+
`wsp.ScriptedLLM.from_transcript("runs/1/transcript.jsonl")` replays a logged run.
|
|
190
|
+
|
|
191
|
+
## Production
|
|
192
|
+
|
|
193
|
+
Budgets (`token_budget`, `cost_budget`, `time_budget`), cost tracking (`pricing=`), `stop_reason`, partial results
|
|
194
|
+
on errors and Ctrl-C, OpenTelemetry tracing, parallel read-only tool calls and replay are covered in
|
|
195
|
+
[docs/production.md](docs/production.md).
|
|
196
|
+
|
|
197
|
+
```python
|
|
198
|
+
from weightspace.otel import OpenTelemetryCallback
|
|
199
|
+
|
|
200
|
+
res = wsp.fix(
|
|
201
|
+
model, task, llm="anthropic:claude-opus-5-5", pricing=(5, 25, 0.5, 6.25), cost_budget=5.0, callbacks=[OpenTelemetryCallback()]
|
|
202
|
+
)
|
|
203
|
+
print(res.stop_reason, res.tokens, res.cost)
|
|
204
|
+
```
|
|
186
205
|
|
|
187
206
|
## CLI
|
|
188
207
|
|
|
@@ -190,8 +209,43 @@ Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script.
|
|
|
190
209
|
weightspace inspect model.safetensors --describe
|
|
191
210
|
weightspace read model.safetensors "What was this trained to do?" --llm openai:gpt-5 --context @notes.md
|
|
192
211
|
weightspace fix model.safetensors "..." --llm anthropic:claude-sonnet-4-5 --save-weights fixed.safetensors --out runs/1
|
|
212
|
+
weightspace fix model.safetensors "..." --llm openai:gpt-5 --probe probes.py:rollout --forward-budget 10 \
|
|
213
|
+
--price 1.25,10 --max-cost 2 --events events.jsonl --json
|
|
214
|
+
weightspace fix model.safetensors "..." --replay runs/1/transcript.jsonl # no tokens spent
|
|
193
215
|
```
|
|
194
216
|
|
|
217
|
+
## Results: Exp. 8 replication
|
|
218
|
+
|
|
219
|
+
We reran the paper's Exp. 8 with this harness: fix Qwen3.5-0.8B's math by editing its weights, no training. Test sets are MATH-500 and GSM8K test, one sample per problem, last `\boxed{}` graded (code in [`experiments/exp8`](experiments/exp8)). ± is one standard error; "cut" is the share of answers that hit the token limit.
|
|
220
|
+
|
|
221
|
+
**Thinking on** (paper protocol, 32k tokens, first 250 problems of each set):
|
|
222
|
+
|
|
223
|
+
| Edit | Editing LLM | Forward passes | Weights changed | MATH-500 | cut | GSM8K | cut |
|
|
224
|
+
|---|---|---|---|---|---|---|---|
|
|
225
|
+
| none (base) | | | 0 | 8.4 ± 1.8 | 92% | 12.4 ± 2.1 | 89% |
|
|
226
|
+
| paper's edit ([checkpoint](https://huggingface.co/RikiyaT/Qwen3.5-0.8B-weight-space-programming)) | gpt-6-astra | 50 | 131,072 | 52.4 ± 3.2 | 34% | 61.6 ± 3.1 | 10% |
|
|
227
|
+
| 4 MLP neurons ×6 | Opus 5.5, thinking high | 7 of 12 | 4,096 | 40.0 ± 3.1 | 60% | 53.2 ± 3.2 | 42% |
|
|
228
|
+
| 4 attention heads ×1.5 | Opus 5.5, thinking off, fast | 7 of 12 | 1,048,576 | 17.2 ± 2.4 | 82% | 34.0 ± 3.0 | 59% |
|
|
229
|
+
| 1 MLP neuron ×0.5 (unfinished run) | gpt-6-astra, thinking high | 12 budget, run lost | 1,024 | not measured yet | | | |
|
|
230
|
+
|
|
231
|
+
**Thinking off** (model card sampling, 8k tokens, full sets: 500 MATH-500, 500 GSM8K):
|
|
232
|
+
|
|
233
|
+
| Edit | MATH-500 | cut | GSM8K | cut |
|
|
234
|
+
|---|---|---|---|---|
|
|
235
|
+
| none (base) | 46.4 ± 2.2 | 31% | 56.6 ± 2.2 | 6% |
|
|
236
|
+
| paper's edit | 42.0 ± 2.2 | 29% | 57.2 ± 2.2 | 3% |
|
|
237
|
+
| Opus, thinking high | 41.8 ± 2.2 | 23% | 54.4 ± 2.2 | 2% |
|
|
238
|
+
| Opus, thinking off | 33.2 ± 2.1 | 20% | 42.0 ± 2.2 | 3% |
|
|
239
|
+
| Astra's neuron (first 250 of each set) | 46.0 ± 3.2 | 28% | 56.4 ± 3.1 | 6% |
|
|
240
|
+
|
|
241
|
+
What this says:
|
|
242
|
+
|
|
243
|
+
- Editing works. Opus found 4 neurons that push down "Wait" and amplified them; 4,096 weights out of 750M took the thinking model from 8.4% / 12.4% to 40.0% / 53.2%.
|
|
244
|
+
- It does not beat the paper. The paper's edit, with 50 forward passes against Opus's 7, is ahead by 8 to 12 points. Our run of the paper's checkpoint (52.4% / 61.6%) is close to what the paper reports (48.1% / 67.1%, with more samples per problem), and our base (8.4% / 12.4%) matches theirs (7.9% / 13.8%).
|
|
245
|
+
- Turning thinking off is a strong baseline. The base model with thinking off scores 46.4% / 56.6%, above Opus's edit, and none of the edits help in non-thinking mode. The edits mostly cut the thinking loops; they don't add math skill. Compare any edit against the thinking-off base.
|
|
246
|
+
- A second Opus session (up to 25 forward passes in total) is not in the table. It ran out of steps without submitting, its best variant (36.6% on practice problems against 30.5% for the ×6 edit) is within noise, and its layer-23 change wasn't saved before the GPU pod died.
|
|
247
|
+
- Training-set gains overstate test gains: Opus measured 8.6% → 50.8% on held-out practice problems, but a fresh practice set gave 30.5%.
|
|
248
|
+
|
|
195
249
|
## Examples
|
|
196
250
|
|
|
197
251
|
- `examples/fix_quadrants.py`: fix a sabotaged classifier (runs offline with a scripted LLM if you omit `--llm`)
|
|
@@ -140,7 +140,23 @@ res = s.run("...")
|
|
|
140
140
|
```
|
|
141
141
|
|
|
142
142
|
Custom tools: `wsp.tool("name", "description", x={"type": "number"})(fn)`.
|
|
143
|
-
Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script
|
|
143
|
+
Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script, and
|
|
144
|
+
`wsp.ScriptedLLM.from_transcript("runs/1/transcript.jsonl")` replays a logged run.
|
|
145
|
+
|
|
146
|
+
## Production
|
|
147
|
+
|
|
148
|
+
Budgets (`token_budget`, `cost_budget`, `time_budget`), cost tracking (`pricing=`), `stop_reason`, partial results
|
|
149
|
+
on errors and Ctrl-C, OpenTelemetry tracing, parallel read-only tool calls and replay are covered in
|
|
150
|
+
[docs/production.md](docs/production.md).
|
|
151
|
+
|
|
152
|
+
```python
|
|
153
|
+
from weightspace.otel import OpenTelemetryCallback
|
|
154
|
+
|
|
155
|
+
res = wsp.fix(
|
|
156
|
+
model, task, llm="anthropic:claude-opus-5-5", pricing=(5, 25, 0.5, 6.25), cost_budget=5.0, callbacks=[OpenTelemetryCallback()]
|
|
157
|
+
)
|
|
158
|
+
print(res.stop_reason, res.tokens, res.cost)
|
|
159
|
+
```
|
|
144
160
|
|
|
145
161
|
## CLI
|
|
146
162
|
|
|
@@ -148,8 +164,43 @@ Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script.
|
|
|
148
164
|
weightspace inspect model.safetensors --describe
|
|
149
165
|
weightspace read model.safetensors "What was this trained to do?" --llm openai:gpt-5 --context @notes.md
|
|
150
166
|
weightspace fix model.safetensors "..." --llm anthropic:claude-sonnet-4-5 --save-weights fixed.safetensors --out runs/1
|
|
167
|
+
weightspace fix model.safetensors "..." --llm openai:gpt-5 --probe probes.py:rollout --forward-budget 10 \
|
|
168
|
+
--price 1.25,10 --max-cost 2 --events events.jsonl --json
|
|
169
|
+
weightspace fix model.safetensors "..." --replay runs/1/transcript.jsonl # no tokens spent
|
|
151
170
|
```
|
|
152
171
|
|
|
172
|
+
## Results: Exp. 8 replication
|
|
173
|
+
|
|
174
|
+
We reran the paper's Exp. 8 with this harness: fix Qwen3.5-0.8B's math by editing its weights, no training. Test sets are MATH-500 and GSM8K test, one sample per problem, last `\boxed{}` graded (code in [`experiments/exp8`](experiments/exp8)). ± is one standard error; "cut" is the share of answers that hit the token limit.
|
|
175
|
+
|
|
176
|
+
**Thinking on** (paper protocol, 32k tokens, first 250 problems of each set):
|
|
177
|
+
|
|
178
|
+
| Edit | Editing LLM | Forward passes | Weights changed | MATH-500 | cut | GSM8K | cut |
|
|
179
|
+
|---|---|---|---|---|---|---|---|
|
|
180
|
+
| none (base) | | | 0 | 8.4 ± 1.8 | 92% | 12.4 ± 2.1 | 89% |
|
|
181
|
+
| paper's edit ([checkpoint](https://huggingface.co/RikiyaT/Qwen3.5-0.8B-weight-space-programming)) | gpt-6-astra | 50 | 131,072 | 52.4 ± 3.2 | 34% | 61.6 ± 3.1 | 10% |
|
|
182
|
+
| 4 MLP neurons ×6 | Opus 5.5, thinking high | 7 of 12 | 4,096 | 40.0 ± 3.1 | 60% | 53.2 ± 3.2 | 42% |
|
|
183
|
+
| 4 attention heads ×1.5 | Opus 5.5, thinking off, fast | 7 of 12 | 1,048,576 | 17.2 ± 2.4 | 82% | 34.0 ± 3.0 | 59% |
|
|
184
|
+
| 1 MLP neuron ×0.5 (unfinished run) | gpt-6-astra, thinking high | 12 budget, run lost | 1,024 | not measured yet | | | |
|
|
185
|
+
|
|
186
|
+
**Thinking off** (model card sampling, 8k tokens, full sets: 500 MATH-500, 500 GSM8K):
|
|
187
|
+
|
|
188
|
+
| Edit | MATH-500 | cut | GSM8K | cut |
|
|
189
|
+
|---|---|---|---|---|
|
|
190
|
+
| none (base) | 46.4 ± 2.2 | 31% | 56.6 ± 2.2 | 6% |
|
|
191
|
+
| paper's edit | 42.0 ± 2.2 | 29% | 57.2 ± 2.2 | 3% |
|
|
192
|
+
| Opus, thinking high | 41.8 ± 2.2 | 23% | 54.4 ± 2.2 | 2% |
|
|
193
|
+
| Opus, thinking off | 33.2 ± 2.1 | 20% | 42.0 ± 2.2 | 3% |
|
|
194
|
+
| Astra's neuron (first 250 of each set) | 46.0 ± 3.2 | 28% | 56.4 ± 3.1 | 6% |
|
|
195
|
+
|
|
196
|
+
What this says:
|
|
197
|
+
|
|
198
|
+
- Editing works. Opus found 4 neurons that push down "Wait" and amplified them; 4,096 weights out of 750M took the thinking model from 8.4% / 12.4% to 40.0% / 53.2%.
|
|
199
|
+
- It does not beat the paper. The paper's edit, with 50 forward passes against Opus's 7, is ahead by 8 to 12 points. Our run of the paper's checkpoint (52.4% / 61.6%) is close to what the paper reports (48.1% / 67.1%, with more samples per problem), and our base (8.4% / 12.4%) matches theirs (7.9% / 13.8%).
|
|
200
|
+
- Turning thinking off is a strong baseline. The base model with thinking off scores 46.4% / 56.6%, above Opus's edit, and none of the edits help in non-thinking mode. The edits mostly cut the thinking loops; they don't add math skill. Compare any edit against the thinking-off base.
|
|
201
|
+
- A second Opus session (up to 25 forward passes in total) is not in the table. It ran out of steps without submitting, its best variant (36.6% on practice problems against 30.5% for the ×6 edit) is within noise, and its layer-23 change wasn't saved before the GPU pod died.
|
|
202
|
+
- Training-set gains overstate test gains: Opus measured 8.6% → 50.8% on held-out practice problems, but a fresh practice set gave 30.5%.
|
|
203
|
+
|
|
153
204
|
## Examples
|
|
154
205
|
|
|
155
206
|
- `examples/fix_quadrants.py`: fix a sabotaged classifier (runs offline with a scripted LLM if you omit `--llm`)
|
|
@@ -0,0 +1,111 @@
|
|
|
1
|
+
"""Time and peak memory of weightspace hot paths on a Qwen-sized synthetic model (CPU).
|
|
2
|
+
|
|
3
|
+
python benchmarks/profile_hotpaths.py [--small] [--json out.json]
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import argparse
|
|
9
|
+
import gc
|
|
10
|
+
import json
|
|
11
|
+
import time
|
|
12
|
+
import tracemalloc
|
|
13
|
+
|
|
14
|
+
import numpy as np
|
|
15
|
+
import torch
|
|
16
|
+
|
|
17
|
+
import weightspace as wsp
|
|
18
|
+
from weightspace.llm import ScriptedLLM, call, reply
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def model(vocab: int, hidden: int, layers: int, inter: int) -> torch.nn.Module:
|
|
22
|
+
m = torch.nn.Module()
|
|
23
|
+
m.embed = torch.nn.Embedding(vocab, hidden)
|
|
24
|
+
for i in range(layers):
|
|
25
|
+
blk = torch.nn.Module()
|
|
26
|
+
for n, shape in {
|
|
27
|
+
"q": (2 * hidden, hidden),
|
|
28
|
+
"k": (hidden // 2, hidden),
|
|
29
|
+
"v": (hidden // 2, hidden),
|
|
30
|
+
"o": (hidden, 2 * hidden),
|
|
31
|
+
"gate": (inter, hidden),
|
|
32
|
+
"up": (inter, hidden),
|
|
33
|
+
"down": (hidden, inter),
|
|
34
|
+
}.items():
|
|
35
|
+
setattr(blk, n, torch.nn.Linear(shape[1], shape[0], bias=False))
|
|
36
|
+
blk.norm = torch.nn.LayerNorm(hidden)
|
|
37
|
+
setattr(m, f"layer{i}", blk)
|
|
38
|
+
return m.to(torch.bfloat16)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
RESULTS: dict[str, dict] = {}
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def bench(label: str, fn, repeat: int = 3):
|
|
45
|
+
gc.collect()
|
|
46
|
+
tracemalloc.start()
|
|
47
|
+
best = float("inf")
|
|
48
|
+
for _ in range(repeat):
|
|
49
|
+
t0 = time.perf_counter()
|
|
50
|
+
out = fn()
|
|
51
|
+
best = min(best, time.perf_counter() - t0)
|
|
52
|
+
peak = tracemalloc.get_traced_memory()[1]
|
|
53
|
+
tracemalloc.stop()
|
|
54
|
+
RESULTS[label] = {"ms": round(best * 1e3, 2), "peak_mb": round(peak / 2**20, 1)}
|
|
55
|
+
print(f"{label:48s} {best * 1e3:10.2f} ms {peak / 2**20:9.1f} MB peak", flush=True)
|
|
56
|
+
return out
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def main() -> None:
|
|
60
|
+
ap = argparse.ArgumentParser()
|
|
61
|
+
ap.add_argument("--small", action="store_true", help="tiny model, for a quick smoke run")
|
|
62
|
+
ap.add_argument("--json")
|
|
63
|
+
a = ap.parse_args()
|
|
64
|
+
torch.manual_seed(0)
|
|
65
|
+
dims = (4096, 256, 2, 896) if a.small else (151936, 1024, 24, 3584)
|
|
66
|
+
m = model(*dims)
|
|
67
|
+
print(f"model: {sum(p.numel() for p in m.parameters()) / 1e6:.0f}M params bf16, torch threads={torch.get_num_threads()}")
|
|
68
|
+
|
|
69
|
+
ws = wsp.WeightSpace.from_torch(m)
|
|
70
|
+
emb, mlp, attn = "embed.weight", "layer0.gate.weight", "layer0.q.weight"
|
|
71
|
+
for n in ws.names():
|
|
72
|
+
ws[n] # materialise float32 mirrors once, outside the timings
|
|
73
|
+
s = wsp.Session(ws, ScriptedLLM([]), mode="fix", max_steps=1)
|
|
74
|
+
t = s.tools
|
|
75
|
+
|
|
76
|
+
bench("list_tensors", lambda: t["list_tensors"].fn(), 20)
|
|
77
|
+
bench("intro prompt", lambda: s._intro("task"), 20)
|
|
78
|
+
bench("describe mlp (cold)", lambda: (ws._describe_cache.clear(), ws.describe(mlp)), 3)
|
|
79
|
+
bench("describe attn (cold)", lambda: (ws._describe_cache.clear(), ws.describe(attn)), 3)
|
|
80
|
+
bench("describe embedding (cold)", lambda: (ws._describe_cache.clear(), ws.describe(emb)), 2)
|
|
81
|
+
bench("read_tensor 64x64", lambda: t["read_tensor"].fn(name=mlp, index="0:64, 0:64"), 10)
|
|
82
|
+
bench("first edit (snapshot) embedding", lambda: (ws._original.pop(emb, None), ws.edit(emb)), 3)
|
|
83
|
+
for n in ws.names("layer*.weight")[:40]:
|
|
84
|
+
ws.edit(n)
|
|
85
|
+
print(f"touched tensors: {len(ws.touched())}, {sum(ws.numel(n) for n in ws.touched()) / 1e6:.0f}M weights")
|
|
86
|
+
bench("write_tensor 8 values (+diff)", lambda: t["write_tensor"].fn(name=emb, values=[0.5] * 8, index="7, 0:8"), 5)
|
|
87
|
+
bench("diff (one change)", ws.diff, 5)
|
|
88
|
+
bench("patch", ws.patch, 3)
|
|
89
|
+
bench("sync (bf16 flush)", ws.sync, 3)
|
|
90
|
+
bench("checkpoint", lambda: ws.checkpoint("c"), 2)
|
|
91
|
+
bench("restore + diff", lambda: t["restore"].fn(label="c"), 2)
|
|
92
|
+
sb = s.sandbox
|
|
93
|
+
bench("sandbox: W[emb][0,0] += 1", lambda: sb.run(f"W['{emb}'][0, 0] += 1"), 10)
|
|
94
|
+
bench("sandbox: trivial expression", lambda: sb.run("1 + 1"), 50)
|
|
95
|
+
|
|
96
|
+
small = wsp.WeightSpace.from_numpy({"w": np.random.randn(64, 64).astype(np.float32)})
|
|
97
|
+
bench("render_space 4k weights", lambda: wsp.render_space(small), 10)
|
|
98
|
+
text = wsp.render_space(small)
|
|
99
|
+
bench("parse_weights 4k weights", lambda: wsp.parse_weights(text), 10)
|
|
100
|
+
|
|
101
|
+
calls = [
|
|
102
|
+
reply(None, call("read_tensor", name=mlp, index=f"{i}, 0:16"), call("describe_tensor", name=attn)) for i in range(30)
|
|
103
|
+
] + [reply(None, call("submit", answer="ok"))]
|
|
104
|
+
bench("agent loop: 30 steps, 60 tool calls", lambda: wsp.read(ws, "q", llm=ScriptedLLM(list(calls)), max_steps=40), 3)
|
|
105
|
+
if a.json:
|
|
106
|
+
with open(a.json, "w") as f:
|
|
107
|
+
json.dump(RESULTS, f, indent=1)
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
if __name__ == "__main__":
|
|
111
|
+
main()
|
|
@@ -3,6 +3,7 @@
|
|
|
3
3
|
- [Concepts](concepts.md): `WeightSpace`, sessions, modes, the sandbox and what the LLM sees.
|
|
4
4
|
- [Probes and budgets](probes.md): how to let the LLM run your network, and how much.
|
|
5
5
|
- [Providers and thinking](providers.md): OpenAI-compatible and Anthropic backends, thinking levels, custom LLMs.
|
|
6
|
+
- [Production](production.md): budgets, cost, stop reasons, events and OpenTelemetry, parallel tools, replay.
|
|
6
7
|
- [Experiments](../experiments/exp8): the Exp. 8 replication on Qwen3.5-0.8B (math with thinking on).
|
|
7
8
|
|
|
8
9
|
The [README](../README.md) has a quick start and the CLI reference.
|
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
# Running in production
|
|
2
|
+
|
|
3
|
+
Everything here is optional. A plain `wsp.fix(model, task, llm=...)` still works the same way.
|
|
4
|
+
|
|
5
|
+
## Results you can act on
|
|
6
|
+
|
|
7
|
+
`Result` tells you why a run stopped and what it cost:
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
res = wsp.fix(model, task, llm="anthropic:claude-opus-5-5", pricing=(5, 25, 0.5, 6.25))
|
|
11
|
+
res.stop_reason # submitted | answered | max_steps | token_budget | cost_budget | time_budget | interrupted | error
|
|
12
|
+
res.ok # True for submitted / answered
|
|
13
|
+
res.tokens # Usage(input_tokens, output_tokens, cached_input_tokens, cache_write_tokens, reasoning_tokens, requests)
|
|
14
|
+
res.cost # USD, when pricing= is given
|
|
15
|
+
res.timings # {"llm": {"calls", "seconds"}, "tool:python": {...}, ...}
|
|
16
|
+
res.run_id # also sent as the OpenAI prompt_cache_key and the trace conversation id
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
`res.usage` still holds the provider's raw counters. `tokens` is the same shape for every provider.
|
|
20
|
+
|
|
21
|
+
## Budgets
|
|
22
|
+
|
|
23
|
+
```python
|
|
24
|
+
wsp.improve(model, goal, llm=..., token_budget=2_000_000, cost_budget=10.0, time_budget=3600, pricing=(1.25, 10))
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
When a budget runs out the LLM gets one last turn to submit, so you still get an answer and the edits so far.
|
|
28
|
+
`pricing` is USD per million tokens: `(input, output[, cached_input, cache_write])`, a dict, or `wsp.Pricing`.
|
|
29
|
+
|
|
30
|
+
## Failures and Ctrl-C
|
|
31
|
+
|
|
32
|
+
- Ctrl-C returns a `Result` with `stop_reason="interrupted"`: transcript, edits and usage up to that point.
|
|
33
|
+
- If the LLM request fails after the SDK's own retries, `wsp.SessionError` is raised. `err.result` has the
|
|
34
|
+
partial result, so nothing is lost.
|
|
35
|
+
- Backends take `max_retries` (default 4) and `timeout` (default 900 s), passed to the official SDKs.
|
|
36
|
+
- Anthropic streams long requests (`stream="auto"`: when `max_tokens` > 21333), which long Opus turns need.
|
|
37
|
+
|
|
38
|
+
## Observability
|
|
39
|
+
|
|
40
|
+
Pass `callbacks=[...]`; each gets every run event (`run_start`, `step_start`, `llm_start`, `llm_end`, `tool_start`,
|
|
41
|
+
`tool_end`, `message`, `run_end`; see `weightspace.events`). A callback is a function or an object with `on_event`.
|
|
42
|
+
|
|
43
|
+
```python
|
|
44
|
+
from weightspace.events import EventLog, LoggingCallback
|
|
45
|
+
from weightspace.otel import OpenTelemetryCallback # pip install "weightspace[otel]"
|
|
46
|
+
|
|
47
|
+
wsp.fix(model, task, llm=..., callbacks=[OpenTelemetryCallback(), LoggingCallback(), EventLog("events.jsonl")])
|
|
48
|
+
```
|
|
49
|
+
|
|
50
|
+
- `OpenTelemetryCallback` emits GenAI semantic-convention spans: `invoke_agent weightspace` per run, `chat <model>`
|
|
51
|
+
per request (token usage, cache reads/writes, finish reason) and `execute_tool <name>` per tool call. Tool
|
|
52
|
+
arguments are only recorded with `capture_content=True`.
|
|
53
|
+
- `LoggingCallback` logs to the `weightspace` logger (milestones at INFO, the rest at DEBUG).
|
|
54
|
+
- `EventLog` / `TranscriptWriter` append JSONL as the run goes, so a crash keeps the log. `log_path=` uses
|
|
55
|
+
`TranscriptWriter`; `verbose=True` uses `ConsoleLogger`.
|
|
56
|
+
|
|
57
|
+
## Parallel tool calls
|
|
58
|
+
|
|
59
|
+
Read-only tools (`list_tensors`, `describe_tensor`, `read_tensor`, and custom tools declared with
|
|
60
|
+
`wsp.tool(..., parallel=True)`) that the LLM asks for in the same turn run on a thread pool
|
|
61
|
+
(`parallel_tools=8` workers). Anything that writes weights, runs Python or spends forward passes runs alone and
|
|
62
|
+
in order, so edits never race. Results always come back in the order they were asked for.
|
|
63
|
+
|
|
64
|
+
## OpenAI
|
|
65
|
+
|
|
66
|
+
- `prompt_cache_key` is set to the run id on api.openai.com, so every step of a run lands on the same cache.
|
|
67
|
+
Pass `prompt_cache_key="..."` to pick your own, or `False` to turn it off.
|
|
68
|
+
- Any other request field goes straight through: `get_llm("openai:gpt-5", parallel_tool_calls=True,
|
|
69
|
+
safety_identifier="user-123")`.
|
|
70
|
+
- `openresponses:<model>` runs the Responses API stateless (`store=False`) and replays encrypted reasoning.
|
|
71
|
+
|
|
72
|
+
## Replay
|
|
73
|
+
|
|
74
|
+
Every run can be replayed offline, with no tokens spent:
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
llm = wsp.ScriptedLLM.from_transcript("runs/1/transcript.jsonl")
|
|
78
|
+
wsp.fix(model, task, llm=llm, probes=[...], forward_budget=10)
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
Tool outputs are recomputed, so replaying on other weights shows how the same program behaves there.
|
|
82
|
+
CLI: `weightspace fix model.safetensors "..." --replay runs/1/transcript.jsonl`.
|
|
83
|
+
|
|
84
|
+
## Sandbox
|
|
85
|
+
|
|
86
|
+
- Code runs under `torch.no_grad()` (`no_grad=True`), so autograd-based training of the weights fails.
|
|
87
|
+
- The timeout works on any thread: SIGALRM on the main thread, a watchdog elsewhere (for example inside a web
|
|
88
|
+
server worker). A single long C call, like a huge matmul, finishes before the timeout fires.
|
|
89
|
+
- It is still not a security boundary. See [SECURITY.md](../SECURITY.md).
|
|
90
|
+
|
|
91
|
+
## Performance
|
|
92
|
+
|
|
93
|
+
Large tensors are processed in 1M-element blocks on a small shared thread pool (NumPy releases the GIL), so
|
|
94
|
+
nothing makes a full float64 copy of a tensor:
|
|
95
|
+
|
|
96
|
+
- `ws.diff()`, `ws.patch()` and checkpoints compare weights bitwise (a NaN that was not edited is not a change).
|
|
97
|
+
The original norm used by `rel_delta` is computed once per tensor.
|
|
98
|
+
- Checkpoints store a sparse delta from the originals, so checkpointing a small edit to an embedding costs
|
|
99
|
+
megabytes, not gigabytes.
|
|
100
|
+
- `describe_tensor` gets exact quantiles from one histogram pass plus a partition of the few bins that hold the
|
|
101
|
+
wanted ranks, instead of sorting the whole tensor. Its text is cached by content, so describing a touched but
|
|
102
|
+
unchanged tensor is free.
|
|
103
|
+
- The spectrum uses an SVD for small matrices and the Gram matrix in the input precision (solved in float64)
|
|
104
|
+
above 256 rows or columns.
|
|
105
|
+
|
|
106
|
+
`python benchmarks/profile_hotpaths.py` times these paths on a 546M-parameter bf16 model (CPU, 8 threads):
|
|
107
|
+
|
|
108
|
+
| | 0.3.0 | now | |
|
|
109
|
+
|---|---|---|---|
|
|
110
|
+
| `describe` embedding 151936 x 1024 | 6,475 ms, 2,376 MB | 1,241 ms, 132 MB | 5.2x |
|
|
111
|
+
| `describe` MLP 3584 x 1024 | 341 ms | 130 ms | 2.6x |
|
|
112
|
+
| `write_tensor` (+ diff over 237M touched weights) | 371 ms | 80 ms | 4.6x |
|
|
113
|
+
| `ws.diff()` | 381 ms | 81 ms | 4.7x |
|
|
114
|
+
| `ws.patch()` | 170 ms | 83 ms | 2.1x |
|
|
115
|
+
| `ws.checkpoint()` | 245 ms, 1,807 MB | 88 ms, 8 MB | 2.8x |
|
|
116
|
+
| `ws.restore()` + diff | 460 ms | 147 ms | 3.1x |
|
|
117
|
+
| first edit of the embedding (snapshot) | 152 ms | 29 ms | 5.3x |
|
|
118
|
+
| agent loop, 30 steps, 60 tool calls | 9,242 ms | 373 ms | 24.8x |
|
|
@@ -19,7 +19,8 @@ Things to know:
|
|
|
19
19
|
- The LLM never gets the model object. It sees tensors through tools (list, describe, read slices, write) and a Python sandbox with `R` (read-only) and `W` (writable) dicts of numpy arrays, plus checkpoint / restore / revert.
|
|
20
20
|
- The only way it can run the network is through probes: functions decorated with `@wsp.probe(cost=1)` whose docstring and arguments become the tool schema. `forward_budget` caps the total cost.
|
|
21
21
|
- Lower level: `wsp.WeightSpace.from_torch / from_hf / from_file / from_numpy`, `ws.diff()`, `ws.save_patch()`, `ws.apply_patch()`, and `wsp.Session(ws, llm, mode=..., probes=..., namespace=..., max_steps=...).run(task)`.
|
|
22
|
-
- `wsp.ScriptedLLM([...])` replays a fixed script for offline tests.
|
|
22
|
+
- `wsp.ScriptedLLM([...])` replays a fixed script for offline tests; `wsp.ScriptedLLM.from_transcript(path)` replays a logged run.
|
|
23
|
+
- Production: `pricing=(in, out)`, `token_budget` / `cost_budget` / `time_budget`, `res.stop_reason`, `res.tokens`, `res.cost`, `callbacks=[...]` (events, logging, `weightspace.otel.OpenTelemetryCallback`). Ctrl-C returns a partial result; API errors raise `wsp.SessionError` with `.result`.
|
|
23
24
|
- The sandbox runs in your process and is not a security boundary.
|
|
24
25
|
|
|
25
26
|
## Docs
|
|
@@ -27,6 +28,7 @@ Things to know:
|
|
|
27
28
|
- [README](https://raw.githubusercontent.com/leofmarciano/weightspace/main/README.md): quick start, thinking table, CLI
|
|
28
29
|
- [Concepts](https://raw.githubusercontent.com/leofmarciano/weightspace/main/docs/concepts.md): WeightSpace, sessions, modes, sandbox, what the LLM sees
|
|
29
30
|
- [Probes and budgets](https://raw.githubusercontent.com/leofmarciano/weightspace/main/docs/probes.md): letting the LLM run your network, and how much
|
|
31
|
+
- [Production](https://raw.githubusercontent.com/leofmarciano/weightspace/main/docs/production.md): budgets, cost, events, OpenTelemetry, parallel tools, replay
|
|
30
32
|
- [Providers and thinking](https://raw.githubusercontent.com/leofmarciano/weightspace/main/docs/providers.md): OpenAI, Open Responses, OpenRouter, vLLM, Anthropic, custom LLMs
|
|
31
33
|
|
|
32
34
|
## Examples
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "weightspace"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.3.0"
|
|
8
8
|
description = "Weight-space programming: let an LLM read and write neural network weights directly."
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.10"
|
|
@@ -30,8 +30,9 @@ torch = ["torch>=2.1"]
|
|
|
30
30
|
hf = ["torch>=2.1", "transformers>=4.40", "safetensors>=0.4"]
|
|
31
31
|
openai = ["openai>=1.40"]
|
|
32
32
|
anthropic = ["anthropic>=0.40"]
|
|
33
|
+
otel = ["opentelemetry-api>=1.20"]
|
|
33
34
|
all = ["torch>=2.1", "transformers>=4.40", "safetensors>=0.4", "openai>=1.40", "anthropic>=0.40"]
|
|
34
|
-
dev = ["pytest>=8", "ruff>=0.5", "mypy>=1.10", "pre-commit>=3"]
|
|
35
|
+
dev = ["opentelemetry-sdk>=1.20", "pytest>=8", "ruff>=0.5", "mypy>=1.10", "pre-commit>=3"]
|
|
35
36
|
|
|
36
37
|
[project.scripts]
|
|
37
38
|
weightspace = "weightspace.cli:main"
|
|
@@ -7,14 +7,16 @@ print(res) # answer + how many weights changed
|
|
|
7
7
|
res.save("runs/fix-1") # transcript, result.json, sparse patch
|
|
8
8
|
"""
|
|
9
9
|
|
|
10
|
-
from .agent import ForwardCall, Result, Session, Tool, tool
|
|
10
|
+
from .agent import ForwardCall, Result, Session, SessionError, Tool, tool
|
|
11
|
+
from .events import ConsoleLogger, Event, EventLog, LoggingCallback, TranscriptWriter
|
|
11
12
|
from .llm import LLM, LLMResponse, ScriptedLLM, ToolCall, get_llm
|
|
12
13
|
from .probes import Probe, probe
|
|
13
14
|
from .render import describe, fmt_array, parse_weights, render_space
|
|
14
15
|
from .space import DiffSummary, NumpyStore, Store, WeightSpace
|
|
15
16
|
from .tasks import as_weightspace, fix, improve, read, write
|
|
17
|
+
from .usage import Pricing, Usage
|
|
16
18
|
|
|
17
|
-
__version__ = "0.
|
|
19
|
+
__version__ = "0.3.0"
|
|
18
20
|
|
|
19
21
|
__all__ = [
|
|
20
22
|
"WeightSpace",
|
|
@@ -22,7 +24,15 @@ __all__ = [
|
|
|
22
24
|
"NumpyStore",
|
|
23
25
|
"DiffSummary",
|
|
24
26
|
"Session",
|
|
27
|
+
"SessionError",
|
|
25
28
|
"Result",
|
|
29
|
+
"Event",
|
|
30
|
+
"ConsoleLogger",
|
|
31
|
+
"LoggingCallback",
|
|
32
|
+
"TranscriptWriter",
|
|
33
|
+
"EventLog",
|
|
34
|
+
"Usage",
|
|
35
|
+
"Pricing",
|
|
26
36
|
"ForwardCall",
|
|
27
37
|
"Tool",
|
|
28
38
|
"tool",
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
"""Block-wise passes over large arrays. NumPy releases the GIL inside its kernels, so big tensors are split
|
|
2
|
+
into blocks that run on a small shared thread pool; small ones stay on the calling thread."""
|
|
3
|
+
|
|
4
|
+
from __future__ import annotations
|
|
5
|
+
|
|
6
|
+
import os
|
|
7
|
+
import threading
|
|
8
|
+
import zlib
|
|
9
|
+
from collections.abc import Callable
|
|
10
|
+
from concurrent.futures import ThreadPoolExecutor
|
|
11
|
+
from typing import TypeVar
|
|
12
|
+
|
|
13
|
+
import numpy as np
|
|
14
|
+
|
|
15
|
+
T = TypeVar("T")
|
|
16
|
+
|
|
17
|
+
CHUNK = 1 << 20 # elements per block
|
|
18
|
+
PARALLEL_MIN = 1 << 22 # below this many elements, one thread is faster than the pool overhead
|
|
19
|
+
WORKERS = max(1, min(8, os.cpu_count() or 1))
|
|
20
|
+
|
|
21
|
+
_pool: ThreadPoolExecutor | None = None
|
|
22
|
+
_pool_lock = threading.Lock()
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def _executor() -> ThreadPoolExecutor:
|
|
26
|
+
global _pool
|
|
27
|
+
if _pool is None:
|
|
28
|
+
with _pool_lock:
|
|
29
|
+
if _pool is None:
|
|
30
|
+
_pool = ThreadPoolExecutor(WORKERS, thread_name_prefix="weightspace-block")
|
|
31
|
+
return _pool
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def map_blocks(fn: Callable[[int, int], T], n: int, block: int | None = None, *, elements: int | None = None) -> list[T]:
|
|
35
|
+
"""`fn(start, stop)` over `[0, n)` in blocks of `block` (default CHUNK), results in order.
|
|
36
|
+
`elements` is the total work size used to decide on threading (defaults to `n`)."""
|
|
37
|
+
block = max(1, block or CHUNK)
|
|
38
|
+
spans = [(i, min(i + block, n)) for i in range(0, n, block)]
|
|
39
|
+
if len(spans) < 2 or WORKERS < 2 or (elements if elements is not None else n) < PARALLEL_MIN:
|
|
40
|
+
return [fn(a, b) for a, b in spans]
|
|
41
|
+
return list(_executor().map(lambda s: fn(*s), spans))
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
_BITS = {2: np.int16, 4: np.int32, 8: np.int64}
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
def bits(a: np.ndarray) -> np.ndarray:
|
|
48
|
+
"""Integer view of a float array, so equality is bitwise (NaN == NaN, a sign flip of 0 counts as an edit)."""
|
|
49
|
+
t = _BITS.get(a.dtype.itemsize) if a.dtype.kind == "f" else None
|
|
50
|
+
return a.view(t) if t is not None else a
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def flat(a: np.ndarray) -> np.ndarray:
|
|
54
|
+
return a.reshape(-1)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def copy(src: np.ndarray) -> np.ndarray:
|
|
58
|
+
"""`src.copy()`, done in parallel blocks for large contiguous arrays."""
|
|
59
|
+
if src.size < PARALLEL_MIN or not src.flags.c_contiguous:
|
|
60
|
+
return src.copy()
|
|
61
|
+
out = np.empty_like(src)
|
|
62
|
+
copy_into(out, src)
|
|
63
|
+
return out
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def copy_into(dst: np.ndarray, src: np.ndarray) -> None:
|
|
67
|
+
"""`dst[...] = src` for same-shape arrays, in parallel blocks when both are large and contiguous."""
|
|
68
|
+
if dst.size < PARALLEL_MIN or not (dst.flags.c_contiguous and src.flags.c_contiguous) or dst.shape != src.shape:
|
|
69
|
+
dst[...] = src
|
|
70
|
+
return
|
|
71
|
+
d, s = flat(dst), flat(src)
|
|
72
|
+
map_blocks(lambda i, j: np.copyto(d[i:j], s[i:j]), d.size)
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
def changed_indices(cur: np.ndarray, orig: np.ndarray) -> np.ndarray:
|
|
76
|
+
"""Flat int64 indices where `cur` and `orig` differ bitwise."""
|
|
77
|
+
a, b = bits(flat(cur)), bits(flat(orig))
|
|
78
|
+
parts = map_blocks(lambda i, j: np.flatnonzero(a[i:j] != b[i:j]) + i, a.size)
|
|
79
|
+
return np.concatenate(parts).astype(np.int64, copy=False) if parts else np.empty(0, np.int64)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def equal(cur: np.ndarray, orig: np.ndarray) -> bool:
|
|
83
|
+
a, b = bits(flat(cur)), bits(flat(orig))
|
|
84
|
+
return all(map_blocks(lambda i, j: bool(np.array_equal(a[i:j], b[i:j])), a.size))
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def sum_squares(a: np.ndarray) -> float:
|
|
88
|
+
x = flat(a)
|
|
89
|
+
|
|
90
|
+
def block(i: int, j: int) -> float:
|
|
91
|
+
c = x[i:j].astype(np.float64)
|
|
92
|
+
return float(np.dot(c, c))
|
|
93
|
+
|
|
94
|
+
return float(sum(map_blocks(block, x.size)))
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def fingerprint(a: np.ndarray) -> tuple:
|
|
98
|
+
"""Cheap content key: shape, dtype and a CRC32 per block."""
|
|
99
|
+
x = flat(a)
|
|
100
|
+
crcs = map_blocks(lambda i, j: zlib.crc32(np.ascontiguousarray(x[i:j]).data), x.size)
|
|
101
|
+
return (a.shape, a.dtype.str, *crcs)
|