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 +139 -0
- scorequant/_binstats.py +88 -0
- scorequant/_chunking.py +34 -0
- scorequant/_execution.py +414 -0
- scorequant/_json.py +55 -0
- scorequant/_typing.py +9 -0
- scorequant/_validation.py +144 -0
- scorequant/api.py +775 -0
- scorequant/artifact.py +452 -0
- scorequant/certify.py +405 -0
- scorequant/components.py +354 -0
- scorequant/config.py +430 -0
- scorequant/criteria.py +130 -0
- scorequant/information.py +673 -0
- scorequant/partition.py +1786 -0
- scorequant/providers.py +465 -0
- scorequant/py.typed +0 -0
- scorequant/quantizers.py +57 -0
- scorequant/ratios.py +347 -0
- scorequant/reports.py +427 -0
- scorequant/result.py +466 -0
- scorequant/solvers/__init__.py +1 -0
- scorequant/solvers/common.py +72 -0
- scorequant/solvers/kmeans.py +234 -0
- scorequant/solvers/scalar.py +142 -0
- scorequant/solvers/soft.py +374 -0
- scorequant/sources.py +362 -0
- scorequant/transforms.py +154 -0
- scorequant/visualization.py +257 -0
- scorequant-0.1.0.dist-info/METADATA +394 -0
- scorequant-0.1.0.dist-info/RECORD +33 -0
- scorequant-0.1.0.dist-info/WHEEL +4 -0
- scorequant-0.1.0.dist-info/licenses/LICENSE +21 -0
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
|
+
]
|
scorequant/_binstats.py
ADDED
|
@@ -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)
|
scorequant/_chunking.py
ADDED
|
@@ -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)))
|
scorequant/_execution.py
ADDED
|
@@ -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
|