flopscope-client 0.4.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.
- flopscope_client-0.4.0/.gitignore +59 -0
- flopscope_client-0.4.0/PKG-INFO +91 -0
- flopscope_client-0.4.0/README.md +81 -0
- flopscope_client-0.4.0/pyproject.toml +30 -0
- flopscope_client-0.4.0/src/flopscope/__init__.py +335 -0
- flopscope_client-0.4.0/src/flopscope/_budget.py +324 -0
- flopscope_client-0.4.0/src/flopscope/_comms_tracker.py +46 -0
- flopscope_client-0.4.0/src/flopscope/_connection.py +187 -0
- flopscope_client-0.4.0/src/flopscope/_constants.py +9 -0
- flopscope_client-0.4.0/src/flopscope/_display.py +283 -0
- flopscope_client-0.4.0/src/flopscope/_getattr.py +52 -0
- flopscope_client-0.4.0/src/flopscope/_math_compat.py +25 -0
- flopscope_client-0.4.0/src/flopscope/_perm_group.py +721 -0
- flopscope_client-0.4.0/src/flopscope/_protocol.py +88 -0
- flopscope_client-0.4.0/src/flopscope/_registry.py +66 -0
- flopscope_client-0.4.0/src/flopscope/_registry_data.py +609 -0
- flopscope_client-0.4.0/src/flopscope/_remote_array.py +691 -0
- flopscope_client-0.4.0/src/flopscope/_weights.py +126 -0
- flopscope_client-0.4.0/src/flopscope/data/__init__.py +1 -0
- flopscope_client-0.4.0/src/flopscope/data/default_weights.json +462 -0
- flopscope_client-0.4.0/src/flopscope/errors.py +107 -0
- flopscope_client-0.4.0/src/flopscope/fft/__init__.py +9 -0
- flopscope_client-0.4.0/src/flopscope/flops.py +140 -0
- flopscope_client-0.4.0/src/flopscope/linalg/__init__.py +68 -0
- flopscope_client-0.4.0/src/flopscope/random/__init__.py +45 -0
- flopscope_client-0.4.0/src/flopscope/stats/__init__.py +82 -0
- flopscope_client-0.4.0/tests/__init__.py +0 -0
- flopscope_client-0.4.0/tests/test_adversarial.py +757 -0
- flopscope_client-0.4.0/tests/test_api_surface.py +437 -0
- flopscope_client-0.4.0/tests/test_budget.py +302 -0
- flopscope_client-0.4.0/tests/test_bugfixes_round2.py +292 -0
- flopscope_client-0.4.0/tests/test_bugfixes_round3.py +507 -0
- flopscope_client-0.4.0/tests/test_comms_tracker.py +70 -0
- flopscope_client-0.4.0/tests/test_connection.py +384 -0
- flopscope_client-0.4.0/tests/test_errors.py +184 -0
- flopscope_client-0.4.0/tests/test_exhaustive.py +1856 -0
- flopscope_client-0.4.0/tests/test_full_integration.py +477 -0
- flopscope_client-0.4.0/tests/test_integration.py +318 -0
- flopscope_client-0.4.0/tests/test_math_compat.py +28 -0
- flopscope_client-0.4.0/tests/test_new_type_encoding.py +58 -0
- flopscope_client-0.4.0/tests/test_protocol_trace.py +598 -0
- flopscope_client-0.4.0/tests/test_remote_array.py +386 -0
- flopscope_client-0.4.0/tests/test_version_handshake.py +106 -0
- flopscope_client-0.4.0/tests/test_weights.py +156 -0
- flopscope_client-0.4.0/uv.lock +389 -0
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
.aicrowd/
|
|
2
|
+
|
|
3
|
+
# Python bytecode & build artifacts
|
|
4
|
+
__pycache__/
|
|
5
|
+
*.pyc
|
|
6
|
+
*.pyo
|
|
7
|
+
*.pyd
|
|
8
|
+
*.egg-info/
|
|
9
|
+
*.egg
|
|
10
|
+
dist/
|
|
11
|
+
build/
|
|
12
|
+
*.whl
|
|
13
|
+
*.so
|
|
14
|
+
|
|
15
|
+
# Virtual environments
|
|
16
|
+
.venv/
|
|
17
|
+
venv/
|
|
18
|
+
env/
|
|
19
|
+
|
|
20
|
+
# IDE & editors
|
|
21
|
+
.vscode/
|
|
22
|
+
.idea/
|
|
23
|
+
*.swp
|
|
24
|
+
*.swo
|
|
25
|
+
*~
|
|
26
|
+
|
|
27
|
+
# Testing & coverage
|
|
28
|
+
.pytest_cache/
|
|
29
|
+
.coverage
|
|
30
|
+
coverage.xml
|
|
31
|
+
htmlcov/
|
|
32
|
+
.tox/
|
|
33
|
+
.nox/
|
|
34
|
+
|
|
35
|
+
# mkdocs generated site
|
|
36
|
+
site/
|
|
37
|
+
|
|
38
|
+
# Website (Fumadocs / Next.js)
|
|
39
|
+
website/node_modules/
|
|
40
|
+
website/.next/
|
|
41
|
+
website/out/
|
|
42
|
+
website/public/llms.txt
|
|
43
|
+
website/public/llms-full.txt
|
|
44
|
+
|
|
45
|
+
# Claude Code session data
|
|
46
|
+
.claude/
|
|
47
|
+
|
|
48
|
+
# Benchmarks (runtime output)
|
|
49
|
+
.benchmarks/
|
|
50
|
+
|
|
51
|
+
# ZMQ sockets (runtime)
|
|
52
|
+
*.sock
|
|
53
|
+
|
|
54
|
+
# OS junk
|
|
55
|
+
.DS_Store
|
|
56
|
+
Thumbs.db
|
|
57
|
+
Desktop.ini
|
|
58
|
+
weights.json
|
|
59
|
+
report.html
|
|
@@ -0,0 +1,91 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: flopscope-client
|
|
3
|
+
Version: 0.4.0
|
|
4
|
+
Summary: Lightweight flopscope client — drop-in replacement with no numpy dependency
|
|
5
|
+
License-Expression: MIT
|
|
6
|
+
Requires-Python: >=3.10
|
|
7
|
+
Requires-Dist: msgpack>=1.0.0
|
|
8
|
+
Requires-Dist: pyzmq>=26.0.0
|
|
9
|
+
Description-Content-Type: text/markdown
|
|
10
|
+
|
|
11
|
+
# flopscope-client
|
|
12
|
+
|
|
13
|
+
[](https://pypi.org/project/flopscope-client/)
|
|
14
|
+
[](https://www.python.org/downloads/)
|
|
15
|
+
[](https://opensource.org/licenses/MIT)
|
|
16
|
+
|
|
17
|
+
Lightweight drop-in replacement for [flopscope](https://pypi.org/project/flopscope/) that proxies all operations to a remote [`flopscope-server`](https://pypi.org/project/flopscope-server/) over ZMQ + msgpack.
|
|
18
|
+
|
|
19
|
+
`flopscope-client` provides the same `import flopscope` Python module as the main `flopscope` distribution, but with **no NumPy dependency** — it forwards every counted operation to a `flopscope-server` process. Use this in constrained environments where you want flopscope's FLOP-counting API but cannot ship numpy + the full library (e.g. sandboxed participant containers in the [ARC Whitebox Estimation Challenge](https://www.alignment.org/blog/will-whitebox-runtime-monitoring-defeat-scheming-models/)).
|
|
20
|
+
|
|
21
|
+
## Install instead of, not alongside
|
|
22
|
+
|
|
23
|
+
`flopscope-client` occupies the same `flopscope` Python import namespace as the main package. The two are **mutually exclusive** — installing both leads to file-overlap in `flopscope/`. Choose one:
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
# Lightweight: client-only, no numpy
|
|
27
|
+
pip install flopscope-client
|
|
28
|
+
|
|
29
|
+
# Heavy: full library on the local machine
|
|
30
|
+
pip install flopscope
|
|
31
|
+
|
|
32
|
+
# Server-side: both flopscope + flopscope-server, pinned together
|
|
33
|
+
pip install "flopscope[server]"
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
## Quick start
|
|
37
|
+
|
|
38
|
+
The client connects to a `flopscope-server` instance specified by the `FLOPSCOPE_SERVER_URL` environment variable:
|
|
39
|
+
|
|
40
|
+
```bash
|
|
41
|
+
export FLOPSCOPE_SERVER_URL=tcp://flopscope-server.example.com:15555
|
|
42
|
+
# or for a local UNIX socket:
|
|
43
|
+
export FLOPSCOPE_SERVER_URL=ipc:///tmp/flopscope.sock
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
Then use flopscope normally — the import path and API are identical to the main distribution:
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
import flopscope as flops
|
|
50
|
+
import flopscope.numpy as fnp
|
|
51
|
+
|
|
52
|
+
with flops.BudgetContext(flop_budget=1_000_000):
|
|
53
|
+
a = fnp.array([1.0, 2.0, 3.0])
|
|
54
|
+
b = fnp.array([4.0, 5.0, 6.0])
|
|
55
|
+
result = fnp.add(a, b) # round-trips to the server, runs there
|
|
56
|
+
flops.budget_summary()
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
On the first request, the client performs a version handshake with the server. A version mismatch raises `ConnectionError` with both versions in the error message — keep `flopscope-server` and `flopscope-client` on the same release.
|
|
60
|
+
|
|
61
|
+
## When to choose this over the main `flopscope` install
|
|
62
|
+
|
|
63
|
+
| Scenario | Install |
|
|
64
|
+
|----------|---------|
|
|
65
|
+
| You want flopscope's API in a process that has full NumPy + scientific stack | `pip install flopscope` |
|
|
66
|
+
| You're shipping a sandboxed container that must not contain numpy / scipy | `pip install flopscope-client` |
|
|
67
|
+
| You're running the server side that hosts computation for many clients | `pip install flopscope-server` (brings `flopscope` along) |
|
|
68
|
+
| You're deploying both sides on one machine, version-pinned | `pip install "flopscope[server]"` |
|
|
69
|
+
|
|
70
|
+
## Architecture overview
|
|
71
|
+
|
|
72
|
+
```
|
|
73
|
+
┌─────────────────────────────┐ ZMQ + msgpack ┌──────────────────────────────┐
|
|
74
|
+
│ flopscope-client process │ ───────────────────────────▶ │ flopscope-server process │
|
|
75
|
+
│ (sandbox; no numpy) │ │ (full flopscope + numpy) │
|
|
76
|
+
│ `import flopscope as flops`│ ◀─────────────────────────── │ Executes ops, tracks budget │
|
|
77
|
+
└─────────────────────────────┘ └──────────────────────────────┘
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
The client serializes each operation (op name, args, kwargs) as msgpack, sends it over the ZMQ REQ/REP socket, and decodes the response. Budget tracking, symmetry-aware FLOP counting, and operation-cost analytics all happen on the server side; the client just relays calls.
|
|
81
|
+
|
|
82
|
+
## Related
|
|
83
|
+
|
|
84
|
+
- [`flopscope`](https://pypi.org/project/flopscope/) — full NumPy-backed library (alternative install)
|
|
85
|
+
- [`flopscope-server`](https://pypi.org/project/flopscope-server/) — the server-side runtime this client connects to
|
|
86
|
+
- [Documentation](https://aicrowd.github.io/flopscope/) — full guides
|
|
87
|
+
- [GitHub](https://github.com/AIcrowd/flopscope) — source, CHANGELOG, contributor guide
|
|
88
|
+
|
|
89
|
+
## License
|
|
90
|
+
|
|
91
|
+
[MIT](https://github.com/AIcrowd/flopscope/blob/main/LICENSE)
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
# flopscope-client
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/flopscope-client/)
|
|
4
|
+
[](https://www.python.org/downloads/)
|
|
5
|
+
[](https://opensource.org/licenses/MIT)
|
|
6
|
+
|
|
7
|
+
Lightweight drop-in replacement for [flopscope](https://pypi.org/project/flopscope/) that proxies all operations to a remote [`flopscope-server`](https://pypi.org/project/flopscope-server/) over ZMQ + msgpack.
|
|
8
|
+
|
|
9
|
+
`flopscope-client` provides the same `import flopscope` Python module as the main `flopscope` distribution, but with **no NumPy dependency** — it forwards every counted operation to a `flopscope-server` process. Use this in constrained environments where you want flopscope's FLOP-counting API but cannot ship numpy + the full library (e.g. sandboxed participant containers in the [ARC Whitebox Estimation Challenge](https://www.alignment.org/blog/will-whitebox-runtime-monitoring-defeat-scheming-models/)).
|
|
10
|
+
|
|
11
|
+
## Install instead of, not alongside
|
|
12
|
+
|
|
13
|
+
`flopscope-client` occupies the same `flopscope` Python import namespace as the main package. The two are **mutually exclusive** — installing both leads to file-overlap in `flopscope/`. Choose one:
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
# Lightweight: client-only, no numpy
|
|
17
|
+
pip install flopscope-client
|
|
18
|
+
|
|
19
|
+
# Heavy: full library on the local machine
|
|
20
|
+
pip install flopscope
|
|
21
|
+
|
|
22
|
+
# Server-side: both flopscope + flopscope-server, pinned together
|
|
23
|
+
pip install "flopscope[server]"
|
|
24
|
+
```
|
|
25
|
+
|
|
26
|
+
## Quick start
|
|
27
|
+
|
|
28
|
+
The client connects to a `flopscope-server` instance specified by the `FLOPSCOPE_SERVER_URL` environment variable:
|
|
29
|
+
|
|
30
|
+
```bash
|
|
31
|
+
export FLOPSCOPE_SERVER_URL=tcp://flopscope-server.example.com:15555
|
|
32
|
+
# or for a local UNIX socket:
|
|
33
|
+
export FLOPSCOPE_SERVER_URL=ipc:///tmp/flopscope.sock
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
Then use flopscope normally — the import path and API are identical to the main distribution:
|
|
37
|
+
|
|
38
|
+
```python
|
|
39
|
+
import flopscope as flops
|
|
40
|
+
import flopscope.numpy as fnp
|
|
41
|
+
|
|
42
|
+
with flops.BudgetContext(flop_budget=1_000_000):
|
|
43
|
+
a = fnp.array([1.0, 2.0, 3.0])
|
|
44
|
+
b = fnp.array([4.0, 5.0, 6.0])
|
|
45
|
+
result = fnp.add(a, b) # round-trips to the server, runs there
|
|
46
|
+
flops.budget_summary()
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
On the first request, the client performs a version handshake with the server. A version mismatch raises `ConnectionError` with both versions in the error message — keep `flopscope-server` and `flopscope-client` on the same release.
|
|
50
|
+
|
|
51
|
+
## When to choose this over the main `flopscope` install
|
|
52
|
+
|
|
53
|
+
| Scenario | Install |
|
|
54
|
+
|----------|---------|
|
|
55
|
+
| You want flopscope's API in a process that has full NumPy + scientific stack | `pip install flopscope` |
|
|
56
|
+
| You're shipping a sandboxed container that must not contain numpy / scipy | `pip install flopscope-client` |
|
|
57
|
+
| You're running the server side that hosts computation for many clients | `pip install flopscope-server` (brings `flopscope` along) |
|
|
58
|
+
| You're deploying both sides on one machine, version-pinned | `pip install "flopscope[server]"` |
|
|
59
|
+
|
|
60
|
+
## Architecture overview
|
|
61
|
+
|
|
62
|
+
```
|
|
63
|
+
┌─────────────────────────────┐ ZMQ + msgpack ┌──────────────────────────────┐
|
|
64
|
+
│ flopscope-client process │ ───────────────────────────▶ │ flopscope-server process │
|
|
65
|
+
│ (sandbox; no numpy) │ │ (full flopscope + numpy) │
|
|
66
|
+
│ `import flopscope as flops`│ ◀─────────────────────────── │ Executes ops, tracks budget │
|
|
67
|
+
└─────────────────────────────┘ └──────────────────────────────┘
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
The client serializes each operation (op name, args, kwargs) as msgpack, sends it over the ZMQ REQ/REP socket, and decodes the response. Budget tracking, symmetry-aware FLOP counting, and operation-cost analytics all happen on the server side; the client just relays calls.
|
|
71
|
+
|
|
72
|
+
## Related
|
|
73
|
+
|
|
74
|
+
- [`flopscope`](https://pypi.org/project/flopscope/) — full NumPy-backed library (alternative install)
|
|
75
|
+
- [`flopscope-server`](https://pypi.org/project/flopscope-server/) — the server-side runtime this client connects to
|
|
76
|
+
- [Documentation](https://aicrowd.github.io/flopscope/) — full guides
|
|
77
|
+
- [GitHub](https://github.com/AIcrowd/flopscope) — source, CHANGELOG, contributor guide
|
|
78
|
+
|
|
79
|
+
## License
|
|
80
|
+
|
|
81
|
+
[MIT](https://github.com/AIcrowd/flopscope/blob/main/LICENSE)
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "flopscope-client"
|
|
7
|
+
version = "0.4.0"
|
|
8
|
+
description = "Lightweight flopscope client — drop-in replacement with no numpy dependency"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
dependencies = [
|
|
13
|
+
"pyzmq>=26.0.0",
|
|
14
|
+
"msgpack>=1.0.0",
|
|
15
|
+
]
|
|
16
|
+
|
|
17
|
+
[tool.hatch.build.targets.wheel]
|
|
18
|
+
packages = ["src/flopscope"]
|
|
19
|
+
|
|
20
|
+
[tool.hatch.build.targets.wheel.force-include]
|
|
21
|
+
"src/flopscope/data/default_weights.json" = "flopscope/data/default_weights.json"
|
|
22
|
+
|
|
23
|
+
[tool.pytest.ini_options]
|
|
24
|
+
pythonpath = ["src"]
|
|
25
|
+
testpaths = ["tests"]
|
|
26
|
+
|
|
27
|
+
[dependency-groups]
|
|
28
|
+
dev = [
|
|
29
|
+
"pytest>=9.0.3",
|
|
30
|
+
]
|
|
@@ -0,0 +1,335 @@
|
|
|
1
|
+
"""flopscope — transparent proxy to a remote flopscope server.
|
|
2
|
+
|
|
3
|
+
This module exposes a numpy-like API where every operation is dispatched
|
|
4
|
+
to a remote server over ZMQ. Participants use it as::
|
|
5
|
+
|
|
6
|
+
import flopscope as flops
|
|
7
|
+
import flopscope.numpy as fnp
|
|
8
|
+
|
|
9
|
+
with flops.BudgetContext(flop_budget=1_000_000) as ctx:
|
|
10
|
+
a = fnp.array([[1.0, 2.0], [3.0, 4.0]])
|
|
11
|
+
b = fnp.zeros((2, 2))
|
|
12
|
+
c = fnp.add(a, b)
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
|
|
17
|
+
import builtins
|
|
18
|
+
import struct
|
|
19
|
+
from typing import Any
|
|
20
|
+
|
|
21
|
+
__version__ = "0.4.0"
|
|
22
|
+
|
|
23
|
+
# ---------------------------------------------------------------------------
|
|
24
|
+
# Errors
|
|
25
|
+
# ---------------------------------------------------------------------------
|
|
26
|
+
|
|
27
|
+
# ---------------------------------------------------------------------------
|
|
28
|
+
# Budget
|
|
29
|
+
# ---------------------------------------------------------------------------
|
|
30
|
+
from flopscope._budget import ( # noqa: E402
|
|
31
|
+
BudgetContext,
|
|
32
|
+
OpRecord,
|
|
33
|
+
budget,
|
|
34
|
+
budget_summary_dict,
|
|
35
|
+
)
|
|
36
|
+
from flopscope._display import budget_live, budget_summary # noqa: E402
|
|
37
|
+
from flopscope._math_compat import e, inf, nan, pi # noqa: E402
|
|
38
|
+
from flopscope._perm_group import SymmetryGroup # noqa: E402
|
|
39
|
+
|
|
40
|
+
# ---------------------------------------------------------------------------
|
|
41
|
+
# Remote types
|
|
42
|
+
# ---------------------------------------------------------------------------
|
|
43
|
+
from flopscope._remote_array import ( # noqa: E402
|
|
44
|
+
_DTYPE_INFO,
|
|
45
|
+
RemoteArray,
|
|
46
|
+
RemoteScalar,
|
|
47
|
+
_encode_arg,
|
|
48
|
+
_result_from_response,
|
|
49
|
+
)
|
|
50
|
+
from flopscope.errors import ( # noqa: E402
|
|
51
|
+
BudgetExhaustedError,
|
|
52
|
+
FlopscopeError,
|
|
53
|
+
FlopscopeServerError,
|
|
54
|
+
FlopscopeWarning,
|
|
55
|
+
NoBudgetContextError,
|
|
56
|
+
SymmetryError,
|
|
57
|
+
)
|
|
58
|
+
|
|
59
|
+
# Alias: ``fnp.ndarray`` refers to the RemoteArray class.
|
|
60
|
+
ndarray = RemoteArray
|
|
61
|
+
|
|
62
|
+
# ---------------------------------------------------------------------------
|
|
63
|
+
# Connection / protocol (private)
|
|
64
|
+
# ---------------------------------------------------------------------------
|
|
65
|
+
|
|
66
|
+
# ---------------------------------------------------------------------------
|
|
67
|
+
# Submodules (imported so ``fnp.linalg``, ``fnp.random``, ``fnp.fft`` work)
|
|
68
|
+
# ---------------------------------------------------------------------------
|
|
69
|
+
from flopscope import (
|
|
70
|
+
fft, # noqa: E402, F401
|
|
71
|
+
flops, # noqa: E402, F401
|
|
72
|
+
linalg, # noqa: E402, F401
|
|
73
|
+
random, # noqa: E402, F401
|
|
74
|
+
stats, # noqa: E402, F401
|
|
75
|
+
)
|
|
76
|
+
from flopscope._connection import get_connection # noqa: E402
|
|
77
|
+
from flopscope._protocol import ( # noqa: E402
|
|
78
|
+
encode_create_from_data,
|
|
79
|
+
encode_request,
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
# ---------------------------------------------------------------------------
|
|
83
|
+
# Registry
|
|
84
|
+
# ---------------------------------------------------------------------------
|
|
85
|
+
from flopscope._registry import ( # noqa: E402
|
|
86
|
+
BLACKLISTED,
|
|
87
|
+
FUNCTION_CATEGORIES,
|
|
88
|
+
get_category,
|
|
89
|
+
is_valid_op,
|
|
90
|
+
iter_proxyable,
|
|
91
|
+
)
|
|
92
|
+
from flopscope._registry_data import FUNCTION_CATEGORIES as _FC # noqa: E402
|
|
93
|
+
|
|
94
|
+
# ---------------------------------------------------------------------------
|
|
95
|
+
# Constants (no server round-trip needed)
|
|
96
|
+
# ---------------------------------------------------------------------------
|
|
97
|
+
|
|
98
|
+
pi: float = pi
|
|
99
|
+
e: float = e
|
|
100
|
+
inf: float = inf
|
|
101
|
+
nan: float = nan
|
|
102
|
+
newaxis = None
|
|
103
|
+
|
|
104
|
+
# ---------------------------------------------------------------------------
|
|
105
|
+
# Dtype strings (mirror numpy dtype names as plain strings)
|
|
106
|
+
# ---------------------------------------------------------------------------
|
|
107
|
+
|
|
108
|
+
float16: str = "float16"
|
|
109
|
+
float32: str = "float32"
|
|
110
|
+
float64: str = "float64"
|
|
111
|
+
int8: str = "int8"
|
|
112
|
+
int16: str = "int16"
|
|
113
|
+
int32: str = "int32"
|
|
114
|
+
int64: str = "int64"
|
|
115
|
+
uint8: str = "uint8"
|
|
116
|
+
bool_: str = "bool"
|
|
117
|
+
complex64: str = "complex64"
|
|
118
|
+
complex128: str = "complex128"
|
|
119
|
+
|
|
120
|
+
# ---------------------------------------------------------------------------
|
|
121
|
+
# Proxy factory
|
|
122
|
+
# ---------------------------------------------------------------------------
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def _make_proxy(op_name: str):
|
|
126
|
+
"""Create a proxy function that dispatches *op_name* to the server."""
|
|
127
|
+
|
|
128
|
+
def proxy(*args: Any, **kwargs: Any):
|
|
129
|
+
conn = get_connection()
|
|
130
|
+
encoded_args = [_encode_arg(a) for a in args]
|
|
131
|
+
encoded_kwargs = {k: _encode_arg(v) for k, v in kwargs.items()}
|
|
132
|
+
resp = conn.send_recv(
|
|
133
|
+
encode_request(op_name, args=encoded_args, kwargs=encoded_kwargs)
|
|
134
|
+
)
|
|
135
|
+
return _result_from_response(resp)
|
|
136
|
+
|
|
137
|
+
proxy.__name__ = op_name
|
|
138
|
+
proxy.__qualname__ = op_name
|
|
139
|
+
return proxy
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
# ---------------------------------------------------------------------------
|
|
143
|
+
# Special-case: array()
|
|
144
|
+
# ---------------------------------------------------------------------------
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
def _flatten(obj):
|
|
148
|
+
"""Recursively flatten a nested list/tuple and return ``(flat, shape)``."""
|
|
149
|
+
if not isinstance(obj, (list, tuple)):
|
|
150
|
+
return [obj], ()
|
|
151
|
+
if len(obj) == 0:
|
|
152
|
+
return [], (0,)
|
|
153
|
+
first_flat, inner_shape = _flatten(obj[0])
|
|
154
|
+
flat = list(first_flat)
|
|
155
|
+
for item in obj[1:]:
|
|
156
|
+
item_flat, item_shape = _flatten(item)
|
|
157
|
+
if item_shape != inner_shape:
|
|
158
|
+
raise ValueError(
|
|
159
|
+
f"Inhomogeneous shape: expected inner shape {inner_shape}, "
|
|
160
|
+
f"got {item_shape}"
|
|
161
|
+
)
|
|
162
|
+
flat.extend(item_flat)
|
|
163
|
+
return flat, (len(obj),) + inner_shape
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
def _infer_dtype(values):
|
|
167
|
+
"""Infer a dtype string from a list of Python scalars."""
|
|
168
|
+
# Use builtins.any/all to avoid collision with the proxy functions
|
|
169
|
+
# that shadow these names at module level.
|
|
170
|
+
_any = builtins.any
|
|
171
|
+
_all = builtins.all
|
|
172
|
+
has_float = _any(isinstance(v, float) for v in values)
|
|
173
|
+
has_complex = _any(isinstance(v, complex) for v in values)
|
|
174
|
+
if has_complex:
|
|
175
|
+
return "complex128"
|
|
176
|
+
if has_float:
|
|
177
|
+
return "float64"
|
|
178
|
+
if _all(isinstance(v, bool) for v in values):
|
|
179
|
+
return "bool"
|
|
180
|
+
if _all(isinstance(v, int) for v in values):
|
|
181
|
+
return "int64"
|
|
182
|
+
return "float64" # mixed or float values
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
def array(object, dtype=None, **kwargs):
|
|
186
|
+
"""Create a remote array from a Python list, tuple, or existing RemoteArray.
|
|
187
|
+
|
|
188
|
+
Parameters
|
|
189
|
+
----------
|
|
190
|
+
object:
|
|
191
|
+
Data to create the array from. May be a nested list/tuple of
|
|
192
|
+
numbers or an existing :class:`RemoteArray`.
|
|
193
|
+
dtype:
|
|
194
|
+
Optional dtype string (e.g. ``"float64"``). Inferred from data
|
|
195
|
+
if not given.
|
|
196
|
+
|
|
197
|
+
Returns
|
|
198
|
+
-------
|
|
199
|
+
RemoteArray
|
|
200
|
+
A new remote array on the server.
|
|
201
|
+
"""
|
|
202
|
+
if isinstance(object, RemoteArray):
|
|
203
|
+
if dtype is None:
|
|
204
|
+
return object
|
|
205
|
+
# dtype cast: dispatch to server
|
|
206
|
+
conn = get_connection()
|
|
207
|
+
resp = conn.send_recv(
|
|
208
|
+
encode_request("astype", args=[{"__handle__": object.handle_id}, dtype])
|
|
209
|
+
)
|
|
210
|
+
return _result_from_response(resp)
|
|
211
|
+
|
|
212
|
+
if isinstance(object, (list, tuple)):
|
|
213
|
+
flat, shape = _flatten(object)
|
|
214
|
+
if not flat:
|
|
215
|
+
# Empty array
|
|
216
|
+
dtype_str = dtype if isinstance(dtype, str) else "float64"
|
|
217
|
+
conn = get_connection()
|
|
218
|
+
resp = conn.send_recv(encode_create_from_data(b"", list(shape), dtype_str))
|
|
219
|
+
return _result_from_response(resp)
|
|
220
|
+
|
|
221
|
+
dtype_str = dtype if isinstance(dtype, str) else (dtype or _infer_dtype(flat))
|
|
222
|
+
info = _DTYPE_INFO.get(dtype_str)
|
|
223
|
+
if info is None:
|
|
224
|
+
raise TypeError(f"Unsupported dtype: {dtype_str!r}")
|
|
225
|
+
fmt_char, _ = info
|
|
226
|
+
|
|
227
|
+
# Complex types: split each value into (real, imag) pairs
|
|
228
|
+
if dtype_str in ("complex64", "complex128"):
|
|
229
|
+
expanded = []
|
|
230
|
+
for v in flat:
|
|
231
|
+
c = complex(v)
|
|
232
|
+
expanded.extend([c.real, c.imag])
|
|
233
|
+
flat = expanded
|
|
234
|
+
fmt_char = "f" if dtype_str == "complex64" else "d"
|
|
235
|
+
data = struct.pack(f"<{len(flat)}{fmt_char}", *flat)
|
|
236
|
+
else:
|
|
237
|
+
data = struct.pack(f"<{len(flat)}{fmt_char}", *flat)
|
|
238
|
+
|
|
239
|
+
conn = get_connection()
|
|
240
|
+
resp = conn.send_recv(encode_create_from_data(data, list(shape), dtype_str))
|
|
241
|
+
return _result_from_response(resp)
|
|
242
|
+
|
|
243
|
+
if isinstance(object, (int, float, complex)):
|
|
244
|
+
# Scalar -> 0-d array
|
|
245
|
+
if isinstance(object, complex) and dtype is None:
|
|
246
|
+
dtype_str = "complex128"
|
|
247
|
+
else:
|
|
248
|
+
dtype_str = dtype if isinstance(dtype, str) else "float64"
|
|
249
|
+
info = _DTYPE_INFO.get(dtype_str)
|
|
250
|
+
if info is None:
|
|
251
|
+
raise TypeError(f"Unsupported dtype: {dtype_str!r}")
|
|
252
|
+
fmt_char, _ = info
|
|
253
|
+
|
|
254
|
+
if dtype_str in ("complex64", "complex128"):
|
|
255
|
+
c = complex(object)
|
|
256
|
+
pack_fmt = "f" if dtype_str == "complex64" else "d"
|
|
257
|
+
data = struct.pack(f"<2{pack_fmt}", c.real, c.imag)
|
|
258
|
+
else:
|
|
259
|
+
data = struct.pack(f"<1{fmt_char}", object)
|
|
260
|
+
conn = get_connection()
|
|
261
|
+
resp = conn.send_recv(encode_create_from_data(data, [], dtype_str))
|
|
262
|
+
return _result_from_response(resp)
|
|
263
|
+
|
|
264
|
+
raise TypeError(
|
|
265
|
+
f"Cannot create array from {type(object).__name__}. "
|
|
266
|
+
f"Expected list, tuple, int, float, or RemoteArray."
|
|
267
|
+
)
|
|
268
|
+
|
|
269
|
+
|
|
270
|
+
# ---------------------------------------------------------------------------
|
|
271
|
+
# Special-case: einsum()
|
|
272
|
+
# ---------------------------------------------------------------------------
|
|
273
|
+
|
|
274
|
+
|
|
275
|
+
def einsum(subscripts, *operands, **kwargs):
|
|
276
|
+
"""Einstein summation on remote arrays.
|
|
277
|
+
|
|
278
|
+
Parameters
|
|
279
|
+
----------
|
|
280
|
+
subscripts:
|
|
281
|
+
Subscript string (e.g. ``"ij,jk->ik"``).
|
|
282
|
+
*operands:
|
|
283
|
+
Input :class:`RemoteArray` objects.
|
|
284
|
+
**kwargs:
|
|
285
|
+
Additional keyword arguments forwarded to the server.
|
|
286
|
+
|
|
287
|
+
Returns
|
|
288
|
+
-------
|
|
289
|
+
RemoteArray
|
|
290
|
+
Result of the einsum operation.
|
|
291
|
+
"""
|
|
292
|
+
conn = get_connection()
|
|
293
|
+
encoded_args = [subscripts] + [_encode_arg(op) for op in operands]
|
|
294
|
+
encoded_kwargs = {k: _encode_arg(v) for k, v in kwargs.items()}
|
|
295
|
+
resp = conn.send_recv(
|
|
296
|
+
encode_request("einsum", args=encoded_args, kwargs=encoded_kwargs)
|
|
297
|
+
)
|
|
298
|
+
return _result_from_response(resp)
|
|
299
|
+
|
|
300
|
+
|
|
301
|
+
# ---------------------------------------------------------------------------
|
|
302
|
+
# Auto-generate proxy functions for all non-blacklisted top-level ops
|
|
303
|
+
# ---------------------------------------------------------------------------
|
|
304
|
+
|
|
305
|
+
# Functions that are special-cased above and should not be overwritten.
|
|
306
|
+
_SPECIAL_CASED = frozenset({"array", "einsum"})
|
|
307
|
+
|
|
308
|
+
# Functions that belong to submodules (contain a dot) are handled by the
|
|
309
|
+
# submodule packages themselves.
|
|
310
|
+
_generated_proxies: list[str] = []
|
|
311
|
+
for _op_name in iter_proxyable():
|
|
312
|
+
if "." in _op_name:
|
|
313
|
+
continue # submodule function
|
|
314
|
+
if _op_name in _SPECIAL_CASED:
|
|
315
|
+
continue
|
|
316
|
+
globals()[_op_name] = _make_proxy(_op_name)
|
|
317
|
+
_generated_proxies.append(_op_name)
|
|
318
|
+
|
|
319
|
+
del _op_name # clean up loop variable
|
|
320
|
+
|
|
321
|
+
|
|
322
|
+
# ---------------------------------------------------------------------------
|
|
323
|
+
# Module-level __getattr__ for blacklisted / unknown names
|
|
324
|
+
# ---------------------------------------------------------------------------
|
|
325
|
+
|
|
326
|
+
# We import the factory but define the function inline so we can also
|
|
327
|
+
# check against names that are already defined in the module namespace.
|
|
328
|
+
|
|
329
|
+
from flopscope._getattr import make_module_getattr as _make_module_getattr # noqa: E402
|
|
330
|
+
|
|
331
|
+
_module_getattr = _make_module_getattr("", "flopscope")
|
|
332
|
+
|
|
333
|
+
|
|
334
|
+
def __getattr__(name: str):
|
|
335
|
+
return _module_getattr(name)
|