scorequant 0.1.0__py3-none-any.whl

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.
scorequant/__init__.py ADDED
@@ -0,0 +1,139 @@
1
+ """Information-preserving hard quantization for statistical inference."""
2
+
3
+ from importlib.metadata import PackageNotFoundError
4
+ from importlib.metadata import version as _distribution_version
5
+
6
+ from .api import fit_quantizer, optimize_partition
7
+ from .artifact import Quantizer
8
+ from .certify import CertificationConfig, certify_partition
9
+ from .components import (
10
+ LinearComponents,
11
+ LinearProblem,
12
+ scores_from_components,
13
+ )
14
+ from .config import (
15
+ DExchangeConfig,
16
+ ExecutionConfig,
17
+ KMeansConfig,
18
+ MahalanobisLloydConfig,
19
+ ScalarDPConfig,
20
+ SoftVoronoiConfig,
21
+ )
22
+ from .criteria import DOptimality, NormalizedTrace, ProfiledDOptimality
23
+ from .information import (
24
+ binned_fisher_information,
25
+ efficient_score_bound,
26
+ efficient_scores,
27
+ fisher_information,
28
+ fractional_fisher_information,
29
+ information_report,
30
+ profiled_information_report,
31
+ )
32
+ from .partition import exchange_stability_report
33
+ from .providers import (
34
+ CentralLogRatioScore,
35
+ DensityRatioScore,
36
+ LinearComponentScore,
37
+ ScoreFunction,
38
+ ScoreProvider,
39
+ )
40
+ from .ratios import (
41
+ IntensityParameterization,
42
+ MixtureParameterization,
43
+ mixture_scores_from_ratios,
44
+ ratio_closure_report,
45
+ ratios_from_posteriors,
46
+ )
47
+ from .reports import RatioClosureReport
48
+ from .result import (
49
+ EfficientScoreBound,
50
+ GeometryReport,
51
+ InformationReport,
52
+ OptimizationTrace,
53
+ PartitionCertificate,
54
+ PartitionResult,
55
+ ProfiledGeometryReport,
56
+ ProfiledInformationReport,
57
+ QuantizerResult,
58
+ StabilityReport,
59
+ )
60
+ from .sources import (
61
+ GaussLegendreConfig,
62
+ IntegrationSource,
63
+ ObservationSample,
64
+ RatioProvenance,
65
+ ScoreProvenance,
66
+ ScoreSample,
67
+ ScoreSchema,
68
+ )
69
+ from .transforms import FisherTransform
70
+ from .visualization import plot_information, plot_optimization, plot_partition, plot_summary
71
+
72
+ #: Installed distribution version. Resolved from package metadata so
73
+ #: ``pyproject.toml`` stays the single source of truth; a source tree that was
74
+ #: never installed reports ``"0.0.0.dev0"`` rather than failing to import.
75
+ try:
76
+ __version__ = _distribution_version("scorequant")
77
+ except PackageNotFoundError: # pragma: no cover - only hit in an uninstalled tree
78
+ __version__ = "0.0.0.dev0"
79
+
80
+ __all__ = [
81
+ "CentralLogRatioScore",
82
+ "CertificationConfig",
83
+ "DExchangeConfig",
84
+ "ExecutionConfig",
85
+ "DOptimality",
86
+ "DensityRatioScore",
87
+ "EfficientScoreBound",
88
+ "FisherTransform",
89
+ "GaussLegendreConfig",
90
+ "GeometryReport",
91
+ "InformationReport",
92
+ "IntegrationSource",
93
+ "IntensityParameterization",
94
+ "KMeansConfig",
95
+ "LinearComponentScore",
96
+ "LinearComponents",
97
+ "LinearProblem",
98
+ "MahalanobisLloydConfig",
99
+ "MixtureParameterization",
100
+ "NormalizedTrace",
101
+ "ObservationSample",
102
+ "OptimizationTrace",
103
+ "PartitionCertificate",
104
+ "PartitionResult",
105
+ "ProfiledDOptimality",
106
+ "ProfiledGeometryReport",
107
+ "ProfiledInformationReport",
108
+ "Quantizer",
109
+ "QuantizerResult",
110
+ "RatioClosureReport",
111
+ "RatioProvenance",
112
+ "ScalarDPConfig",
113
+ "ScoreFunction",
114
+ "ScoreProvenance",
115
+ "ScoreProvider",
116
+ "ScoreSample",
117
+ "ScoreSchema",
118
+ "SoftVoronoiConfig",
119
+ "StabilityReport",
120
+ "binned_fisher_information",
121
+ "certify_partition",
122
+ "efficient_score_bound",
123
+ "efficient_scores",
124
+ "exchange_stability_report",
125
+ "fisher_information",
126
+ "fit_quantizer",
127
+ "fractional_fisher_information",
128
+ "information_report",
129
+ "mixture_scores_from_ratios",
130
+ "optimize_partition",
131
+ "plot_information",
132
+ "plot_optimization",
133
+ "plot_partition",
134
+ "plot_summary",
135
+ "profiled_information_report",
136
+ "ratio_closure_report",
137
+ "ratios_from_posteriors",
138
+ "scores_from_components",
139
+ ]
@@ -0,0 +1,88 @@
1
+ """Shared weighted per-bin scatter-add statistics.
2
+
3
+ Every hard-label solver and diagnostic in ScoreQuant repeats the same three
4
+ steps on a fresh set of integer bin labels: scatter-add the row weights into
5
+ per-bin occupancy, scatter-add the weighted rows into per-bin sums, and divide
6
+ to get per-bin means. :func:`scatter_bin_statistics` is that one
7
+ implementation.
8
+
9
+ Empty-bin policy
10
+ -----------------
11
+ A bin with zero total weight divides its (also zero) sum by one instead of by
12
+ its own zero weight, so its reported mean is exactly ``0.0`` rather than
13
+ ``NaN``: ``jnp.where(bin_weights > 0, bin_weights, 1)``. This is the
14
+ ``where``-guard used everywhere in ScoreQuant that a hard-label mean can
15
+ legitimately be evaluated on a labeling with an empty declared cell (for
16
+ example, an externally supplied labeling passed to
17
+ :func:`scorequant.information.information_report`).
18
+
19
+ Callers for whom an empty cell must instead be a caller error (the exact
20
+ exchange engine, which must never propose a state with an unoccupied
21
+ requested cell) raise using the returned ``weights`` array themselves; the
22
+ guard still keeps the arithmetic that produces ``means`` well-defined before
23
+ that check runs, so the two concerns stay independent.
24
+
25
+ This is deliberately not used for the differentiable soft-responsibility path
26
+ in ``quantizers.soft_voronoi`` (``_soft_fisher``), which floors its occupancy
27
+ with ``jnp.maximum(occupancy, tiny)`` instead. That path's occupancy is a sum
28
+ of continuous softmax responsibilities that is never exactly zero but can be
29
+ arbitrarily small, and it is differentiated through by ``jax.grad``: flooring
30
+ keeps the gradient of the mean finite and well-scaled as the occupancy
31
+ shrinks, while the ``where``-guard's discontinuous branch does not have a
32
+ useful gradient at the switch point. The two guards solve different problems
33
+ and are not interchangeable.
34
+ """
35
+
36
+ from __future__ import annotations
37
+
38
+ from dataclasses import dataclass
39
+
40
+ from ._execution import scatter_add
41
+ from ._execution import xp as jnp
42
+
43
+
44
+ @dataclass(frozen=True, slots=True)
45
+ class BinStatistics:
46
+ """Per-bin weighted occupancy, value sums, and empty-bin-safe means."""
47
+
48
+ weights: jnp.ndarray
49
+ sums: jnp.ndarray
50
+ means: jnp.ndarray
51
+
52
+
53
+ def scatter_bin_statistics(
54
+ labels: jnp.ndarray,
55
+ weights: jnp.ndarray,
56
+ values: jnp.ndarray,
57
+ n_bins: int,
58
+ ) -> BinStatistics:
59
+ """Accumulate weighted occupancy, value sums, and safe per-bin means.
60
+
61
+ Parameters
62
+ ----------
63
+ labels
64
+ Integer bin label for every row, with shape ``[N]`` and values in
65
+ ``[0, n_bins)``.
66
+ weights
67
+ Nonnegative row weights with shape ``[N]``.
68
+ values
69
+ Row values to accumulate, with shape ``[N, D]``.
70
+ n_bins
71
+ Total number of bins, including bins no label selects.
72
+
73
+ Returns
74
+ -------
75
+ BinStatistics
76
+ ``weights`` and ``sums`` with shape ``[n_bins]`` and ``[n_bins, D]``,
77
+ and ``means = sums / weights`` with the empty-bin policy documented on
78
+ this module.
79
+ """
80
+ bin_weights = scatter_add(jnp.zeros(n_bins, dtype=weights.dtype), labels, weights)
81
+ bin_sums = scatter_add(
82
+ jnp.zeros((n_bins, values.shape[1]), dtype=values.dtype),
83
+ labels,
84
+ weights[:, None] * values,
85
+ )
86
+ safe_weights = jnp.where(bin_weights > 0, bin_weights, 1)
87
+ means = bin_sums / safe_weights[:, None]
88
+ return BinStatistics(weights=bin_weights, sums=bin_sums, means=means)
@@ -0,0 +1,34 @@
1
+ """Shared memory-bounded row-chunking budget for dense distance kernels.
2
+
3
+ Several private kernels evaluate one ``[chunk_rows, n_bins]`` (or the
4
+ Mahalanobis ``[chunk_rows, n_bins, rank]`` residual behind it) distance table
5
+ per row chunk instead of materializing the full ``[n_rows, n_bins, rank]``
6
+ tensor at once. ``assignment_chunk_rows`` is the one place that budget is
7
+ computed, so ``partition.py``, ``quantizers.py``, and ``result.py`` size their
8
+ chunks identically instead of drifting apart under independent edits.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import numpy as np
14
+ import numpy.typing as npt
15
+
16
+ # One chunk's dense temporaries (the base distance table plus its
17
+ # per-dimension residual and einsum working set) are held inside this many
18
+ # bytes, independent of total row count.
19
+ WORKING_SET_BYTES = 64 * 1024 * 1024
20
+
21
+
22
+ def assignment_chunk_rows(dtype: npt.DTypeLike, n_rows: int, n_bins: int, rank: int) -> int:
23
+ """Return how many rows one memory-bounded assignment chunk holds.
24
+
25
+ ``n_bins * (rank + 4) + 4 * rank`` accounts for the ``[chunk, n_bins]``
26
+ distance table, the ``[chunk, n_bins, rank]`` residual and einsum
27
+ temporaries a Mahalanobis assignment needs, and a fixed allowance for the
28
+ Euclidean case's smaller working set; it is a deliberately generous
29
+ single formula shared by every chunked assignment kernel rather than a
30
+ per-kernel estimate.
31
+ """
32
+ item_size = np.dtype(dtype).itemsize
33
+ values_per_row = n_bins * (rank + 4) + 4 * rank
34
+ return max(1, min(n_rows, WORKING_SET_BYTES // max(item_size * values_per_row, 1)))
@@ -0,0 +1,414 @@
1
+ """Private execution resolution and backend primitive adapters.
2
+
3
+ Backend selection is scoped with :mod:`contextvars`, so nested public calls
4
+ inherit one resolved runtime without process-global mutable state. Numerical
5
+ modules import ``xp`` instead of a concrete array namespace; backend-name
6
+ conditionals stay in this module.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from collections.abc import Callable, Iterator
12
+ from contextlib import contextmanager, nullcontext
13
+ from contextvars import ContextVar
14
+ from dataclasses import dataclass, fields, is_dataclass, replace
15
+ from functools import cache, wraps
16
+ from types import ModuleType
17
+ from typing import TYPE_CHECKING, Protocol, cast, overload
18
+
19
+ import numpy as np
20
+
21
+ from ._typing import ArrayLike
22
+ from .config import DeviceKind, ExecutionConfig
23
+
24
+
25
+ @dataclass(frozen=True, slots=True)
26
+ class _Runtime:
27
+ config: ExecutionConfig
28
+ namespace: ModuleType
29
+ device: object | None
30
+
31
+
32
+ _ACTIVE_RUNTIME: ContextVar[_Runtime | None] = ContextVar("scorequant_active_runtime", default=None)
33
+
34
+
35
+ def _resolve_runtime(config: ExecutionConfig) -> _Runtime:
36
+ if config.backend == "numpy":
37
+ return _Runtime(
38
+ config=replace(config, device="cpu"),
39
+ namespace=np,
40
+ device=None,
41
+ )
42
+
43
+ try:
44
+ import jax
45
+ import jax.numpy as jax_numpy
46
+ except ImportError as error: # pragma: no cover - exercised in isolated import smoke test
47
+ raise RuntimeError(
48
+ "the JAX backend is unavailable; install ScoreQuant on CPython with its default "
49
+ "dependencies or select ExecutionConfig(backend='numpy')"
50
+ ) from error
51
+
52
+ if config.precision == "float64" and not bool(jax.config.x64_enabled):
53
+ raise RuntimeError(
54
+ "float64 was requested for the JAX backend, but JAX_ENABLE_X64 is disabled; "
55
+ "enable X64 before importing JAX or choose precision='float32'"
56
+ )
57
+ requested_platform = None if config.device == "default" else config.device
58
+ devices = jax.devices(requested_platform)
59
+ if not devices:
60
+ raise RuntimeError(f"no JAX {config.device!r} device is available")
61
+ device = devices[0]
62
+ resolved = replace(config, device=cast(DeviceKind, device.platform))
63
+ return _Runtime(config=resolved, namespace=jax_numpy, device=device)
64
+
65
+
66
+ @contextmanager
67
+ def use_execution(execution: ExecutionConfig | None) -> Iterator[ExecutionConfig]:
68
+ """Enter a validated execution scope, inheriting an existing scope if omitted."""
69
+ active = _ACTIVE_RUNTIME.get()
70
+ if execution is None and active is not None:
71
+ yield active.config
72
+ return
73
+
74
+ runtime = _resolve_runtime(execution or ExecutionConfig())
75
+ token = _ACTIVE_RUNTIME.set(runtime)
76
+ device_scope = nullcontext()
77
+ if runtime.config.backend == "jax":
78
+ import jax
79
+
80
+ device_scope = jax.default_device(runtime.device)
81
+ try:
82
+ with device_scope:
83
+ yield runtime.config
84
+ finally:
85
+ _ACTIVE_RUNTIME.reset(token)
86
+
87
+
88
+ def current_execution() -> ExecutionConfig:
89
+ """Return the active resolved execution, defaulting lazily to JAX."""
90
+ runtime = _ACTIVE_RUNTIME.get()
91
+ if runtime is None:
92
+ return _resolve_runtime(ExecutionConfig()).config
93
+ return runtime.config
94
+
95
+
96
+ def _current_runtime() -> _Runtime:
97
+ runtime = _ACTIVE_RUNTIME.get()
98
+ return _resolve_runtime(ExecutionConfig()) if runtime is None else runtime
99
+
100
+
101
+ class _NamespaceProxy:
102
+ """Resolve NumPy-compatible namespace attributes from the active runtime."""
103
+
104
+ def __getattr__(self, name: str) -> object:
105
+ return getattr(_current_runtime().namespace, name)
106
+
107
+
108
+ if TYPE_CHECKING:
109
+ # Shared kernels use NumPy's typed API. At runtime the proxy resolves the
110
+ # equivalent primitive from the active NumPy or JAX namespace.
111
+ import numpy as xp
112
+ else:
113
+ xp = _NamespaceProxy()
114
+
115
+
116
+ def apply_precision(array: ArrayLike) -> np.ndarray:
117
+ """Convert a numerical input according to the active precision policy."""
118
+ runtime = _current_runtime()
119
+ values = runtime.namespace.asarray(array)
120
+ dtype = values.dtype
121
+ if runtime.config.precision == "float32":
122
+ return cast(np.ndarray, values.astype(runtime.namespace.float32))
123
+ if runtime.config.precision == "float64":
124
+ return cast(np.ndarray, values.astype(runtime.namespace.float64))
125
+ if not runtime.namespace.issubdtype(dtype, runtime.namespace.inexact):
126
+ return cast(np.ndarray, values.astype(runtime.namespace.float32))
127
+ if dtype in (runtime.namespace.float16, getattr(runtime.namespace, "bfloat16", object())):
128
+ return cast(np.ndarray, values.astype(runtime.namespace.float32))
129
+ return cast(np.ndarray, values)
130
+
131
+
132
+ class _FunctionalUpdate(Protocol):
133
+ def add(self, values: object) -> object: ...
134
+
135
+ def set(self, values: object) -> object: ...
136
+
137
+
138
+ class _ArrayAtIndexer(Protocol):
139
+ def __getitem__(self, indices: object) -> _FunctionalUpdate: ...
140
+
141
+
142
+ class _FunctionalArray(Protocol):
143
+ @property
144
+ def at(self) -> _ArrayAtIndexer: ...
145
+
146
+
147
+ def scatter_add(base: np.ndarray, indices: object, values: object) -> np.ndarray:
148
+ """Return ``base`` with repeated-index additions applied functionally."""
149
+ if _current_runtime().config.backend == "jax":
150
+ return cast(np.ndarray, cast(_FunctionalArray, base).at[indices].add(values))
151
+ result = np.array(base, copy=True)
152
+ np.add.at(result, np.asarray(indices), np.asarray(values))
153
+ return result
154
+
155
+
156
+ def scatter_set(base: np.ndarray, indices: object, values: object) -> np.ndarray:
157
+ """Return ``base`` with selected entries replaced functionally."""
158
+ if _current_runtime().config.backend == "jax":
159
+ return cast(np.ndarray, cast(_FunctionalArray, base).at[indices].set(values))
160
+ result = np.array(base, copy=True)
161
+ result[np.asarray(indices)] = np.asarray(values)
162
+ return result
163
+
164
+
165
+ def scatter_block_add(base: np.ndarray, row_indices: object, values: object) -> np.ndarray:
166
+ """Add a square block selected by one index vector."""
167
+ if _current_runtime().config.backend == "jax":
168
+ namespace = _current_runtime().namespace
169
+ selection = namespace.ix_(row_indices, row_indices)
170
+ return cast(np.ndarray, cast(_FunctionalArray, base).at[selection].add(values))
171
+ result = np.array(base, copy=True)
172
+ selection = np.ix_(np.asarray(row_indices), np.asarray(row_indices))
173
+ result[selection] += np.asarray(values)
174
+ return result
175
+
176
+
177
+ type RandomSeed = int | tuple[int, int]
178
+
179
+
180
+ def _key_seed(key: object) -> tuple[int, int]:
181
+ values = np.asarray(key, dtype=np.uint32)
182
+ return int(values[0]), int(values[1])
183
+
184
+
185
+ def split_seeds(seed: RandomSeed, count: int) -> tuple[RandomSeed, ...]:
186
+ """Return deterministic independent seeds for the active backend."""
187
+ if _current_runtime().config.backend == "jax":
188
+ import jax
189
+
190
+ key = (
191
+ jax.random.PRNGKey(seed) if isinstance(seed, int) else xp.asarray(seed, dtype=xp.uint32)
192
+ )
193
+ keys = jax.random.split(key, count)
194
+ return tuple(_key_seed(key) for key in keys)
195
+ sequence = np.random.SeedSequence(seed)
196
+ return tuple(
197
+ int(child.generate_state(1, dtype=np.uint32)[0]) for child in sequence.spawn(count)
198
+ )
199
+
200
+
201
+ def weighted_choice(seed: RandomSeed, probabilities: object) -> int:
202
+ """Draw one categorical index with a backend-local deterministic seed."""
203
+ values = np.asarray(probabilities, dtype=np.float64)
204
+ values = values / np.sum(values)
205
+ if _current_runtime().config.backend == "jax":
206
+ import jax
207
+
208
+ key = (
209
+ jax.random.PRNGKey(seed) if isinstance(seed, int) else xp.asarray(seed, dtype=xp.uint32)
210
+ )
211
+ return int(np.asarray(jax.random.choice(key, values.shape[0], p=xp.asarray(values))))
212
+ return int(np.random.default_rng(seed).choice(values.shape[0], p=values))
213
+
214
+
215
+ def random_permutation(seed: int, size: int) -> np.ndarray:
216
+ """Return one backend-local deterministic permutation."""
217
+ if _current_runtime().config.backend == "jax":
218
+ import jax
219
+
220
+ return cast(np.ndarray, jax.random.permutation(jax.random.PRNGKey(seed), size))
221
+ return np.random.default_rng(seed).permutation(size).astype(np.int32)
222
+
223
+
224
+ class _OptimizerTransformation(Protocol):
225
+ def init(self, parameters: np.ndarray) -> object: ...
226
+
227
+ def update(
228
+ self,
229
+ gradients: np.ndarray,
230
+ state: object,
231
+ parameters: np.ndarray,
232
+ ) -> tuple[np.ndarray, object]: ...
233
+
234
+
235
+ @dataclass(slots=True)
236
+ class AdamState:
237
+ """Opaque backend-owned Adam state shared by the solver orchestration."""
238
+
239
+ learning_rate: float
240
+ gradient_clip: float
241
+ step: int = 0
242
+ first_moment: np.ndarray | None = None
243
+ second_moment: np.ndarray | None = None
244
+ transformation: _OptimizerTransformation | None = None
245
+ backend_state: object | None = None
246
+
247
+
248
+ def create_adam(
249
+ parameters: np.ndarray,
250
+ *,
251
+ learning_rate: float,
252
+ gradient_clip: float,
253
+ ) -> AdamState:
254
+ """Create the active backend's private Adam state."""
255
+ if _current_runtime().config.backend == "jax":
256
+ import optax
257
+
258
+ transformation = cast(
259
+ _OptimizerTransformation,
260
+ optax.chain(
261
+ optax.clip_by_global_norm(gradient_clip),
262
+ optax.adam(learning_rate, b1=0.9, b2=0.999, eps=1e-8),
263
+ ),
264
+ )
265
+ return AdamState(
266
+ learning_rate=learning_rate,
267
+ gradient_clip=gradient_clip,
268
+ transformation=transformation,
269
+ backend_state=transformation.init(parameters),
270
+ )
271
+ values = np.asarray(parameters)
272
+ return AdamState(
273
+ learning_rate=learning_rate,
274
+ gradient_clip=gradient_clip,
275
+ first_moment=np.zeros_like(values),
276
+ second_moment=np.zeros_like(values),
277
+ )
278
+
279
+
280
+ def adam_update(
281
+ parameters: np.ndarray,
282
+ gradients: np.ndarray,
283
+ state: AdamState,
284
+ ) -> tuple[np.ndarray, AdamState, float]:
285
+ """Apply clipping and one bias-corrected Adam update."""
286
+ gradient_norm = float(np.linalg.norm(np.asarray(gradients)))
287
+ if _current_runtime().config.backend == "jax":
288
+ import optax
289
+
290
+ transformation = state.transformation
291
+ if transformation is None:
292
+ raise RuntimeError("JAX Adam state is missing its transformation")
293
+ updates, backend_state = transformation.update(
294
+ gradients,
295
+ state.backend_state,
296
+ parameters,
297
+ )
298
+ state.backend_state = backend_state
299
+ state.step += 1
300
+ return cast(np.ndarray, optax.apply_updates(parameters, updates)), state, gradient_norm
301
+
302
+ gradient = np.asarray(gradients)
303
+ if gradient_norm > state.gradient_clip:
304
+ gradient = gradient * (state.gradient_clip / gradient_norm)
305
+ if state.first_moment is None or state.second_moment is None:
306
+ raise RuntimeError("NumPy Adam state is uninitialized")
307
+ first = 0.9 * np.asarray(state.first_moment) + 0.1 * gradient
308
+ second = 0.999 * np.asarray(state.second_moment) + 0.001 * gradient**2
309
+ state.step += 1
310
+ first_hat = first / (1 - 0.9**state.step)
311
+ second_hat = second / (1 - 0.999**state.step)
312
+ updated = np.asarray(parameters) - state.learning_rate * first_hat / (
313
+ np.sqrt(second_hat) + 1e-8
314
+ )
315
+ state.first_moment = first
316
+ state.second_moment = second
317
+ return updated, state, gradient_norm
318
+
319
+
320
+ @cache
321
+ def _compiled[**P, R](function: Callable[P, R], static_argnames: tuple[str, ...]) -> Callable[P, R]:
322
+ import jax
323
+
324
+ return jax.jit(function, static_argnames=static_argnames)
325
+
326
+
327
+ @overload
328
+ def backend_jit[**P, R](function: Callable[P, R], /) -> Callable[P, R]: ...
329
+
330
+
331
+ @overload
332
+ def backend_jit[**P, R](
333
+ function: None = None, /, *, static_argnames: tuple[str, ...]
334
+ ) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
335
+
336
+
337
+ def backend_jit[**P, R](
338
+ function: Callable[P, R] | None = None,
339
+ /,
340
+ *,
341
+ static_argnames: tuple[str, ...] = (),
342
+ ) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
343
+ """Compile one shared kernel under JAX and call it directly under NumPy."""
344
+
345
+ def decorate(target: Callable[P, R]) -> Callable[P, R]:
346
+ @wraps(target)
347
+ def dispatched(*args: P.args, **kwargs: P.kwargs) -> R:
348
+ if _current_runtime().config.backend == "jax":
349
+ return _compiled(target, static_argnames)(*args, **kwargs)
350
+ return target(*args, **kwargs)
351
+
352
+ return dispatched
353
+
354
+ return decorate(function) if function is not None else decorate
355
+
356
+
357
+ def execution_scope[**P, R](function: Callable[P, R]) -> Callable[P, R]:
358
+ """Resolve the ``execution`` keyword for one public numerical operation."""
359
+
360
+ @wraps(function)
361
+ def wrapped(*args: P.args, **kwargs: P.kwargs) -> R:
362
+ execution = kwargs.get("execution")
363
+ if execution is not None and not isinstance(execution, ExecutionConfig):
364
+ raise TypeError("execution must be an ExecutionConfig or None")
365
+ with use_execution(execution):
366
+ return function(*args, **kwargs)
367
+
368
+ return wrapped
369
+
370
+
371
+ def canonical_array(value: object) -> np.ndarray:
372
+ """Copy one backend array into the stable NumPy API representation."""
373
+ return np.asarray(value).copy()
374
+
375
+
376
+ def backend_array(value: np.ndarray) -> np.ndarray:
377
+ """Place one canonical NumPy array on the active backend.
378
+
379
+ The inverse of :func:`canonical_array`, for kernels re-entered from the
380
+ NumPy arrays a public result stores. ``jax.device_put`` is a transfer
381
+ rather than a staged primitive, so unlike ``jnp.asarray`` it costs no XLA
382
+ compilation on the first call for a given shape. The device is left
383
+ implicit: :func:`use_execution` is already inside
384
+ ``jax.default_device(runtime.device)``, and naming the device here commits
385
+ the array through a resharding path two orders of magnitude slower.
386
+ """
387
+ runtime = _current_runtime()
388
+ if runtime.config.backend == "jax":
389
+ import jax
390
+
391
+ return cast(np.ndarray, jax.device_put(value))
392
+ return np.asarray(value)
393
+
394
+
395
+ def canonicalize_public[T](value: T) -> T:
396
+ """Recursively replace arrays in a public dataclass tree with NumPy arrays."""
397
+ if isinstance(value, np.ndarray):
398
+ return value.copy() # type: ignore[return-value]
399
+ value_type = type(value)
400
+ if value_type.__module__.startswith("jax") or value_type.__module__.startswith("numpy"):
401
+ if hasattr(value, "shape") and hasattr(value, "dtype"):
402
+ return cast(T, np.asarray(value).copy())
403
+ if is_dataclass(value) and not isinstance(value, type):
404
+ for record_field in fields(value):
405
+ current = getattr(value, record_field.name)
406
+ canonical = canonicalize_public(current)
407
+ if canonical is not current:
408
+ object.__setattr__(value, record_field.name, canonical)
409
+ return value
410
+ if isinstance(value, tuple):
411
+ return cast(T, tuple(canonicalize_public(item) for item in value))
412
+ if isinstance(value, list):
413
+ return cast(T, [canonicalize_public(item) for item in value])
414
+ return value