From ae6a0db4638732949832dddabf6bd01565dbea76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Tue, 6 Oct 2026 23:42:43 +0200 Subject: [PATCH 01/17] chore: adopt Conventional Commits for commit messages --- .pre-commit-config.yaml | 12 ++++++++++++ AGENTS.md | 10 ++++++++-- docs/development.md | 26 ++++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b220c65..b5753f4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,3 +1,9 @@ +# default_stages keeps the file-filtered hooks above from running at the +# commit-msg stage, where pre-commit would otherwise feed them +# .git/COMMIT_EDITMSG; default_install_hook_types makes +# `pre-commit install` install both hook types. +default_install_hook_types: [pre-commit, commit-msg] +default_stages: [pre-commit] repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 @@ -30,3 +36,9 @@ repos: language: system files: ^(pyproject\.toml|uv\.lock)$ pass_filenames: false + - repo: https://github.com/compilerla/conventional-pre-commit + rev: v4.4.0 + hooks: + - id: conventional-pre-commit + stages: [commit-msg] + args: [feat, fix, docs, style, refactor, perf, test, build, ci, chore, revert] diff --git a/AGENTS.md b/AGENTS.md index 5b2408c..b33e80c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -95,5 +95,11 @@ then link it for Claude: divergences, and API. `README.md` only links to it. A user-visible change updates the matching page in the same commit; moving modules updates `docs/architecture.md`. -- Commit subjects are `area: summary` in lowercase, e.g. `codegen: ...`, - `lower: ...`, `tests: ...`, `docs: ...`, `ci: ...`, `build: ...`. +- Commit subjects follow [Conventional Commits 1.0.0] + (https://conventionalcommits.org): `type(scope): summary`, lowercase. + Types: `build`, `chore`, `ci`, `docs`, `feat`, `fix`, `perf`, + `refactor`, `revert`, `style`, `test`. The scope is optional but + recommended, from the package areas: `codegen`, `lower`, `optimize`, + `frontend`, `ir`, `runtime`, `api`, `pipeline`, `fuzz`, `agents`. + Mark breaking changes with `!` after the type/scope, or with a + `BREAKING CHANGE` footer. diff --git a/docs/development.md b/docs/development.md index 1183618..8204fb0 100644 --- a/docs/development.md +++ b/docs/development.md @@ -106,6 +106,32 @@ git push origin main --tags or the lock is stale; tagging stays a manual step so a re-run of the target can never move a tag. +## Commit messages + +Commit subjects follow [Conventional Commits 1.0.0] +(https://conventionalcommits.org): `type(scope): summary`, lowercase. +Types: `build`, `chore`, `ci`, `docs`, `feat`, `fix`, `perf`, `refactor`, +`revert`, `style`, `test`. The scope is optional but recommended, from +the package areas: `codegen`, `lower`, `optimize`, `frontend`, `ir`, +`runtime`, `api`, `pipeline`, `fuzz`, `agents`. Mark breaking changes +with `!` after the type/scope, or with a `BREAKING CHANGE` footer. A +`commit-msg` hook (`conventional-pre-commit`) enforces the format. + +The history before the switch used `area: summary` subjects; map them +to the new format like this: + +| Old subject | Conventional equivalent | +|---|---| +| `codegen: x` | `feat(codegen): x` or `fix(codegen): x` | +| `agents: x` | `chore(agents): x` | +| `tests: x` | `test: x` (type, no scope) | +| `docs: x` | `docs: x` (type, no scope) | +| `ci: x` | `ci: x` (type, no scope) | +| `build: x` | `build: x` (type, no scope) | + +Existing clones must re-run `make install` once to pick up the new +commit-msg hook. + ## Coding agents Instructions and skills for coding agents are harness-agnostic; each From 80993e5033e1b67de31642148a84cbab8ee0f788 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 00:48:54 +0200 Subject: [PATCH 02/17] test: split the inspectability suite by topic --- docs/development.md | 2 +- tests/behavior/test_docstrings.py | 115 ++++++++++ tests/behavior/test_inspectability.py | 131 +++++++++++ tests/behavior/test_lambdas.py | 14 ++ tests/behavior/test_memoization.py | 91 ++++++++ tests/behavior/test_public_api.py | 9 + tests/test_inspectability.py | 301 -------------------------- 7 files changed, 361 insertions(+), 302 deletions(-) create mode 100644 tests/behavior/test_docstrings.py create mode 100644 tests/behavior/test_inspectability.py create mode 100644 tests/behavior/test_lambdas.py create mode 100644 tests/behavior/test_memoization.py create mode 100644 tests/behavior/test_public_api.py delete mode 100644 tests/test_inspectability.py diff --git a/docs/development.md b/docs/development.md index 8204fb0..4364604 100644 --- a/docs/development.md +++ b/docs/development.md @@ -53,7 +53,7 @@ plain `pytest` skips nothing but keeps benchmarks disabled | backends (array-api-strict) | `tests/backends/` | | differential (Hypothesis, edge values) | `tests/differential/` | | golden source snapshots | `tests/golden/` (`cases/`, ast.dump equality) | -| inspectability (getsource, linecache) | `tests/test_inspectability.py` | +| inspectability (getsource, linecache, memoization, docstrings) | `tests/behavior/test_inspectability.py` etc. | | grammar fuzzer | `src/array_vectorize/fuzz.py` + CI seeds | | fuzz smoke | `tests/test_fuzz_smoke.py` | | performance gate + benchmarks | `tests/test_perf.py`, `tests/test_bench.py` | diff --git a/tests/behavior/test_docstrings.py b/tests/behavior/test_docstrings.py new file mode 100644 index 0000000..a091bb8 --- /dev/null +++ b/tests/behavior/test_docstrings.py @@ -0,0 +1,115 @@ +"""Docstrings on vectorized functions: summary prefix, source round trip, Notes merging.""" + +from __future__ import annotations + +import inspect + +import pytest + +from array_vectorize import vectorize +from array_vectorize.codegen.docstring import extract_scalar_source + + +def relu_fn(x: float) -> float: + if x < 0: + return 0.0 + return x + + +def sigmoid_like(x: float, alpha: float = 1.0) -> float: + """A smooth sigmoid-like squashing function. + + Maps ``x`` into ``(0, 1)`` using ``alpha`` as the steepness. + """ + if x < 0: + return 0.0 + return 1.0 / (1.0 + alpha * x) + + +def numpy_doc_fn(x: float) -> float: + """Square with a NumPy-style docstring. + + Parameters + ---------- + x : float + input value + + Returns + ------- + float + the square + """ + return x * x + + +def notes_doc_fn(x: float) -> float: + """Clamp nonnegative values. + + Notes: + Pre-existing note. + """ + if x < 0: + return 0.0 + return x + + +def marker_quoting_fn(x: float) -> float: + """Quotes the array-vectorize marker in its own docs. + + Vectorized by array-vectorize from this scalar original:: + + not the real source + """ + return x * 3.0 + + +def test_name_and_docstring() -> None: + vec = vectorize(relu_fn) + assert vec.__name__ == "relu_fn" + doc = inspect.getdoc(vec) or "" + assert doc.splitlines()[0] == "(array-vectorized) no docstring on the scalar original." + assert "def relu_fn(x:" in doc + assert "if x < 0:" in doc + + +def test_docstring_prefixed_summary() -> None: + vec = vectorize(sigmoid_like) + doc = inspect.getdoc(vec) + assert doc is not None + assert doc.splitlines()[0] == "(array-vectorized) A smooth sigmoid-like squashing function." + assert "Maps ``x`` into ``(0, 1)`` using ``alpha`` as the steepness." in doc + + +@pytest.mark.parametrize( + "fn", [relu_fn, sigmoid_like, numpy_doc_fn, notes_doc_fn], ids=lambda f: f.__name__ +) +def test_docstring_source_round_trip(fn) -> None: + # the Notes section embeds the scalar source verbatim; cleandoc of the + # generated docstring recovers it exactly + vec = vectorize(fn) + doc = inspect.getdoc(vec) + assert doc is not None + assert extract_scalar_source(doc) == inspect.getsource(fn).strip() + + +def test_docstring_numpy_style_notes() -> None: + # a NumPy-style original gets a NumPy-style Notes section + vec = vectorize(numpy_doc_fn) + doc = inspect.getdoc(vec) + assert "Notes\n-----" in (doc or "") + + +def test_docstring_existing_notes_merged() -> None: + vec = vectorize(notes_doc_fn) + doc = inspect.getdoc(vec) or "" + # exactly one header line (the copy inside the embedded source is indented) + assert [line for line in doc.splitlines() if line == "Notes:"] == ["Notes:"] + assert "Pre-existing note." in doc + + +def test_docstring_marker_quoting_degrades_to_legacy() -> None: + # the original's docs quote the marker line, so the combined round + # trip cannot be guaranteed: the legacy source-only docstring is + # emitted instead (cleandoc-exact, like the pre-redesign format) + vec = vectorize(marker_quoting_fn) + assert inspect.getdoc(vec) == inspect.getsource(marker_quoting_fn).strip() diff --git a/tests/behavior/test_inspectability.py b/tests/behavior/test_inspectability.py new file mode 100644 index 0000000..d53d5bb --- /dev/null +++ b/tests/behavior/test_inspectability.py @@ -0,0 +1,131 @@ +"""Inspectability of vectorized functions: getsource, linecache, tracebacks, help().""" + +from __future__ import annotations + +import inspect +import linecache +import pickle +import traceback + +import pytest + +from array_vectorize import VectorizationError, get_source, vectorize + + +def relu_fn(x: float) -> float: + if x < 0: + return 0.0 + return x + + +def sigmoid_like(x: float, alpha: float = 1.0) -> float: + """A smooth sigmoid-like squashing function. + + Maps ``x`` into ``(0, 1)`` using ``alpha`` as the steepness. + """ + if x < 0: + return 0.0 + return 1.0 / (1.0 + alpha * x) + + +def test_source_attribute() -> None: + vec = vectorize(relu_fn) + assert isinstance(vec.source, str) + assert "xp.where" in vec.source + + +def test_getsource_returns_generated_not_original() -> None: + vec = vectorize(relu_fn) + src = inspect.getsource(vec) + assert "array_namespace" in src + assert "def relu_fn_vec" in src + assert "def relu_fn(x:" not in src.split('"""')[0] # only inside docstring + + +def test_getsourcelines() -> None: + vec = vectorize(relu_fn) + lines, lineno = inspect.getsourcelines(vec) + assert any("array_namespace" in line for line in lines) + assert lineno >= 1 + + +def test_linecache_survives_checkcache() -> None: + vec = vectorize(relu_fn) + filename = vec.__code__.co_filename + assert filename == "" + linecache.checkcache(filename) + assert linecache.cache[filename][2] # still registered + assert "array_namespace" in "".join(linecache.getlines(filename)) + + +def test_traceback_renders_generated_lines() -> None: + vec = vectorize(relu_fn) + try: + vec(1.0) # no array among args -> namespace resolution fails + except TypeError: + text = traceback.format_exc() + assert "" in text + else: + raise AssertionError("expected TypeError for all-scalar call") + + +def test_wrapped_points_to_original() -> None: + vec = vectorize(relu_fn) + # __wrapped__ is deliberately not set (it would hide the generated source + # from inspect.getsourcelines, which unwraps unconditionally); the + # original is reachable via the marker and drives the signature. + assert vec._vectorized_original is relu_fn + assert inspect.signature(vec) == inspect.signature(relu_fn) + + +def test_module_and_qualname_from_original() -> None: + # module attribution makes pydoc render "Help on function f in module m" + # and is safe: getsource/getsourcelines use co_filename + linecache + vec = vectorize(relu_fn) + assert vec.__module__ == relu_fn.__module__ + assert vec.__qualname__ == relu_fn.__qualname__ + + +def test_pickle_fails_loudly_never_aliases() -> None: + # __module__/__qualname__ match the original, so pickle resolves the + # name and fails the identity check — it cannot silently serialize + # (and unpickle to) the scalar original + vec = vectorize(relu_fn) + with pytest.raises(pickle.PicklingError): + pickle.dumps(vec) + + +def test_signature_preserved() -> None: + def with_defaults(x: float, scale: float = 2.0, *, bias: float = 0.5) -> float: + return x * scale + bias + + vec = vectorize(with_defaults) + sig = inspect.signature(vec) + assert list(sig.parameters) == ["x", "scale", "bias"] + assert sig.parameters["scale"].default == 2.0 + assert sig.parameters["bias"].default == 0.5 + assert sig.parameters["bias"].kind is inspect.Parameter.KEYWORD_ONLY + + +def test_get_source_helper() -> None: + vec = vectorize(relu_fn) + assert get_source(vec) == vec.source + with pytest.raises(VectorizationError, match="not a vectorized function"): + get_source(relu_fn) + + +def test_help_renders() -> None: + # help() must not raise; it shows the module header, the prefixed + # summary, and the original source in the Notes section + vec = vectorize(sigmoid_like) + import io + from contextlib import redirect_stdout + + buf = io.StringIO() + with redirect_stdout(buf): + help(vec) + out = buf.getvalue() + assert f"Help on function sigmoid_like in module {sigmoid_like.__module__}:" in out + assert "(array-vectorized) A smooth sigmoid-like squashing function." in out + assert "Notes:" in out + assert "def sigmoid_like(x: float, alpha: float = 1.0) -> float:" in out diff --git a/tests/behavior/test_lambdas.py b/tests/behavior/test_lambdas.py new file mode 100644 index 0000000..bc2226e --- /dev/null +++ b/tests/behavior/test_lambdas.py @@ -0,0 +1,14 @@ +"""Vectorizing lambdas: identification, naming, values.""" + +from __future__ import annotations + +import numpy as np + +from array_vectorize import vectorize + + +def test_lambda_vectorize() -> None: + vec = vectorize(lambda x: x * 2.0) + got = vec(np.asarray([1.0, 2.0])) + assert np.allclose(got, [2.0, 4.0]) + assert "def lambda_vec" in vec.source diff --git a/tests/behavior/test_memoization.py b/tests/behavior/test_memoization.py new file mode 100644 index 0000000..fa957e3 --- /dev/null +++ b/tests/behavior/test_memoization.py @@ -0,0 +1,91 @@ +"""Memoization of vectorize(): canonical cache, backrefs, idempotent revectorization.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from array_vectorize import vectorize + + +def relu_fn(x: float) -> float: + if x < 0: + return 0.0 + return x + + +def sigmoid_like(x: float, alpha: float = 1.0) -> float: + """A smooth sigmoid-like squashing function. + + Maps ``x`` into ``(0, 1)`` using ``alpha`` as the steepness. + """ + if x < 0: + return 0.0 + return 1.0 / (1.0 + alpha * x) + + +def helper_inner_fn(y: float) -> float: + return y * y + 1.0 + + +def helper_outer_fn(x: float) -> float: + return helper_inner_fn(x) + helper_inner_fn(x * 2.0) + + +def test_backref_on_canonical_vectorization() -> None: + vec = vectorize(sigmoid_like) + assert sigmoid_like.__array_vectorized__ is vec + + +def test_vectorize_memoized_per_original() -> None: + assert vectorize(sigmoid_like) is vectorize(sigmoid_like) + assert vectorize(sigmoid_like) is sigmoid_like.__array_vectorized__ + + +def test_revectorization_returns_same_object() -> None: + vec = vectorize(sigmoid_like) + assert vectorize(vec) is vec + + +def test_verify_runs_on_cache_hit() -> None: + args = (np.asarray([-1.0, 0.5, 2.0]),) + vectorize(sigmoid_like, verify=args) + vectorize(sigmoid_like, verify=args) + + +def test_backref_not_overwritten_by_variants() -> None: + vec = vectorize(sigmoid_like) + vectorize(sigmoid_like, protect_domains=True) + assert sigmoid_like.__array_vectorized__ is vec + vectorize(sigmoid_like, namespace=np) + assert sigmoid_like.__array_vectorized__ is vec + + +def test_backref_on_fallback_wrapper() -> None: + def bad(x: float) -> float: + while x > 0: + x = x - 1 + return x + + with pytest.warns(UserWarning, match="falling back"): + vec = vectorize(bad, fallback=True) + assert bad.__array_vectorized__ is vec + + +def test_helpers_do_not_get_backref() -> None: + # D7 helper compilations are internal: the user's helper function + # must not grow an __array_vectorized__ marker + vectorize(helper_outer_fn) + assert not hasattr(helper_inner_fn, "__array_vectorized__") + assert helper_outer_fn.__array_vectorized__ is not None + + +def test_idempotent_revectorization() -> None: + vec1 = vectorize(relu_fn) + vec2 = vectorize(vec1) + x = np.asarray([-1.0, 2.0]) + assert np.allclose(vec1(x), vec2(x)) + assert "xp.where" in vec2.source + # canonical calls are memoized per scalar original: re-vectorizing + # returns the very same object + assert vec2 is vec1 diff --git a/tests/behavior/test_public_api.py b/tests/behavior/test_public_api.py new file mode 100644 index 0000000..6800762 --- /dev/null +++ b/tests/behavior/test_public_api.py @@ -0,0 +1,9 @@ +"""Public API surface of the array_vectorize package.""" + +from __future__ import annotations + +import array_vectorize + + +def test_package_exposes_version() -> None: + assert array_vectorize.__version__ diff --git a/tests/test_inspectability.py b/tests/test_inspectability.py deleted file mode 100644 index bd5c089..0000000 --- a/tests/test_inspectability.py +++ /dev/null @@ -1,301 +0,0 @@ -"""Inspectability tests: getsource, linecache, tracebacks, attributes.""" - -from __future__ import annotations - -import inspect -import linecache -import pickle -import traceback - -import numpy as np -import pytest - -import array_vectorize -from array_vectorize import VectorizationError, get_source, vectorize -from array_vectorize.codegen.docstring import extract_scalar_source - - -def relu_fn(x: float) -> float: - if x < 0: - return 0.0 - return x - - -def sigmoid_like(x: float, alpha: float = 1.0) -> float: - """A smooth sigmoid-like squashing function. - - Maps ``x`` into ``(0, 1)`` using ``alpha`` as the steepness. - """ - if x < 0: - return 0.0 - return 1.0 / (1.0 + alpha * x) - - -def numpy_doc_fn(x: float) -> float: - """Square with a NumPy-style docstring. - - Parameters - ---------- - x : float - input value - - Returns - ------- - float - the square - """ - return x * x - - -def notes_doc_fn(x: float) -> float: - """Clamp nonnegative values. - - Notes: - Pre-existing note. - """ - if x < 0: - return 0.0 - return x - - -def marker_quoting_fn(x: float) -> float: - """Quotes the array-vectorize marker in its own docs. - - Vectorized by array-vectorize from this scalar original:: - - not the real source - """ - return x * 3.0 - - -def helper_inner_fn(y: float) -> float: - return y * y + 1.0 - - -def helper_outer_fn(x: float) -> float: - return helper_inner_fn(x) + helper_inner_fn(x * 2.0) - - -def test_source_attribute() -> None: - vec = vectorize(relu_fn) - assert isinstance(vec.source, str) - assert "xp.where" in vec.source - - -def test_getsource_returns_generated_not_original() -> None: - vec = vectorize(relu_fn) - src = inspect.getsource(vec) - assert "array_namespace" in src - assert "def relu_fn_vec" in src - assert "def relu_fn(x:" not in src.split('"""')[0] # only inside docstring - - -def test_getsourcelines() -> None: - vec = vectorize(relu_fn) - lines, lineno = inspect.getsourcelines(vec) - assert any("array_namespace" in line for line in lines) - assert lineno >= 1 - - -def test_linecache_survives_checkcache() -> None: - vec = vectorize(relu_fn) - filename = vec.__code__.co_filename - assert filename == "" - linecache.checkcache(filename) - assert linecache.cache[filename][2] # still registered - assert "array_namespace" in "".join(linecache.getlines(filename)) - - -def test_traceback_renders_generated_lines() -> None: - vec = vectorize(relu_fn) - try: - vec(1.0) # no array among args -> namespace resolution fails - except TypeError: - text = traceback.format_exc() - assert "" in text - else: - raise AssertionError("expected TypeError for all-scalar call") - - -def test_wrapped_points_to_original() -> None: - vec = vectorize(relu_fn) - # __wrapped__ is deliberately not set (it would hide the generated source - # from inspect.getsourcelines, which unwraps unconditionally); the - # original is reachable via the marker and drives the signature. - assert vec._vectorized_original is relu_fn - assert inspect.signature(vec) == inspect.signature(relu_fn) - - -def test_module_and_qualname_from_original() -> None: - # module attribution makes pydoc render "Help on function f in module m" - # and is safe: getsource/getsourcelines use co_filename + linecache - vec = vectorize(relu_fn) - assert vec.__module__ == relu_fn.__module__ - assert vec.__qualname__ == relu_fn.__qualname__ - - -def test_pickle_fails_loudly_never_aliases() -> None: - # __module__/__qualname__ match the original, so pickle resolves the - # name and fails the identity check — it cannot silently serialize - # (and unpickle to) the scalar original - vec = vectorize(relu_fn) - with pytest.raises(pickle.PicklingError): - pickle.dumps(vec) - - -def test_backref_on_canonical_vectorization() -> None: - vec = vectorize(sigmoid_like) - assert sigmoid_like.__array_vectorized__ is vec - - -def test_vectorize_memoized_per_original() -> None: - assert vectorize(sigmoid_like) is vectorize(sigmoid_like) - assert vectorize(sigmoid_like) is sigmoid_like.__array_vectorized__ - - -def test_revectorization_returns_same_object() -> None: - vec = vectorize(sigmoid_like) - assert vectorize(vec) is vec - - -def test_verify_runs_on_cache_hit() -> None: - args = (np.asarray([-1.0, 0.5, 2.0]),) - vectorize(sigmoid_like, verify=args) - vectorize(sigmoid_like, verify=args) - - -def test_backref_not_overwritten_by_variants() -> None: - vec = vectorize(sigmoid_like) - vectorize(sigmoid_like, protect_domains=True) - assert sigmoid_like.__array_vectorized__ is vec - vectorize(sigmoid_like, namespace=np) - assert sigmoid_like.__array_vectorized__ is vec - - -def test_backref_on_fallback_wrapper() -> None: - def bad(x: float) -> float: - while x > 0: - x = x - 1 - return x - - with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(bad, fallback=True) - assert bad.__array_vectorized__ is vec - - -def test_helpers_do_not_get_backref() -> None: - # D7 helper compilations are internal: the user's helper function - # must not grow an __array_vectorized__ marker - vectorize(helper_outer_fn) - assert not hasattr(helper_inner_fn, "__array_vectorized__") - assert helper_outer_fn.__array_vectorized__ is not None - - -def test_signature_preserved() -> None: - def with_defaults(x: float, scale: float = 2.0, *, bias: float = 0.5) -> float: - return x * scale + bias - - vec = vectorize(with_defaults) - sig = inspect.signature(vec) - assert list(sig.parameters) == ["x", "scale", "bias"] - assert sig.parameters["scale"].default == 2.0 - assert sig.parameters["bias"].default == 0.5 - assert sig.parameters["bias"].kind is inspect.Parameter.KEYWORD_ONLY - - -def test_name_and_docstring() -> None: - vec = vectorize(relu_fn) - assert vec.__name__ == "relu_fn" - doc = inspect.getdoc(vec) or "" - assert doc.splitlines()[0] == "(array-vectorized) no docstring on the scalar original." - assert "def relu_fn(x:" in doc - assert "if x < 0:" in doc - - -def test_docstring_prefixed_summary() -> None: - vec = vectorize(sigmoid_like) - doc = inspect.getdoc(vec) - assert doc is not None - assert doc.splitlines()[0] == "(array-vectorized) A smooth sigmoid-like squashing function." - assert "Maps ``x`` into ``(0, 1)`` using ``alpha`` as the steepness." in doc - - -@pytest.mark.parametrize( - "fn", [relu_fn, sigmoid_like, numpy_doc_fn, notes_doc_fn], ids=lambda f: f.__name__ -) -def test_docstring_source_round_trip(fn) -> None: - # the Notes section embeds the scalar source verbatim; cleandoc of the - # generated docstring recovers it exactly - vec = vectorize(fn) - doc = inspect.getdoc(vec) - assert doc is not None - assert extract_scalar_source(doc) == inspect.getsource(fn).strip() - - -def test_docstring_numpy_style_notes() -> None: - # a NumPy-style original gets a NumPy-style Notes section - vec = vectorize(numpy_doc_fn) - doc = inspect.getdoc(vec) - assert "Notes\n-----" in (doc or "") - - -def test_docstring_existing_notes_merged() -> None: - vec = vectorize(notes_doc_fn) - doc = inspect.getdoc(vec) or "" - # exactly one header line (the copy inside the embedded source is indented) - assert [line for line in doc.splitlines() if line == "Notes:"] == ["Notes:"] - assert "Pre-existing note." in doc - - -def test_docstring_marker_quoting_degrades_to_legacy() -> None: - # the original's docs quote the marker line, so the combined round - # trip cannot be guaranteed: the legacy source-only docstring is - # emitted instead (cleandoc-exact, like the pre-redesign format) - vec = vectorize(marker_quoting_fn) - assert inspect.getdoc(vec) == inspect.getsource(marker_quoting_fn).strip() - - -def test_idempotent_revectorization() -> None: - vec1 = vectorize(relu_fn) - vec2 = vectorize(vec1) - x = np.asarray([-1.0, 2.0]) - assert np.allclose(vec1(x), vec2(x)) - assert "xp.where" in vec2.source - # canonical calls are memoized per scalar original: re-vectorizing - # returns the very same object - assert vec2 is vec1 - - -def test_get_source_helper() -> None: - vec = vectorize(relu_fn) - assert get_source(vec) == vec.source - with pytest.raises(VectorizationError, match="not a vectorized function"): - get_source(relu_fn) - - -def test_lambda_vectorize() -> None: - vec = vectorize(lambda x: x * 2.0) - got = vec(np.asarray([1.0, 2.0])) - assert np.allclose(got, [2.0, 4.0]) - assert "def lambda_vec" in vec.source - - -def test_help_renders() -> None: - # help() must not raise; it shows the module header, the prefixed - # summary, and the original source in the Notes section - vec = vectorize(sigmoid_like) - import io - from contextlib import redirect_stdout - - buf = io.StringIO() - with redirect_stdout(buf): - help(vec) - out = buf.getvalue() - assert f"Help on function sigmoid_like in module {sigmoid_like.__module__}:" in out - assert "(array-vectorized) A smooth sigmoid-like squashing function." in out - assert "Notes:" in out - assert "def sigmoid_like(x: float, alpha: float = 1.0) -> float:" in out - - -def test_package_exposes_version() -> None: - assert array_vectorize.__version__ From 90458ec356fc3bfd8272194d5c710564dd00f325 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 00:54:09 +0200 Subject: [PATCH 03/17] test: split closures and helpers by section --- tests/behavior/test_closures.py | 44 ++++ tests/behavior/test_closures_and_helpers.py | 259 -------------------- tests/behavior/test_cse.py | 28 +++ tests/behavior/test_helpers.py | 114 +++++++++ tests/behavior/test_lambdas.py | 42 ++++ tests/behavior/test_name_allocation.py | 59 +++++ 6 files changed, 287 insertions(+), 259 deletions(-) create mode 100644 tests/behavior/test_closures.py delete mode 100644 tests/behavior/test_closures_and_helpers.py create mode 100644 tests/behavior/test_cse.py create mode 100644 tests/behavior/test_helpers.py create mode 100644 tests/behavior/test_name_allocation.py diff --git a/tests/behavior/test_closures.py b/tests/behavior/test_closures.py new file mode 100644 index 0000000..d5988a0 --- /dev/null +++ b/tests/behavior/test_closures.py @@ -0,0 +1,44 @@ +"""Closures over scalars and arrays, keyword-only defaults.""" + +from __future__ import annotations + +import corpus +import numpy as np +from support import make_module, vfn + +from array_vectorize import vectorize + +X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) + + +# ------------------------------------------------------- closures and defaults + + +def test_closure_scalar() -> None: + assert np.allclose(vfn("closure_scalar")(X), X * corpus.SCALE) + + +def test_closure_array() -> None: + got = vfn("closure_array")(np.asarray([1.0, 1.0, 1.0])) + assert np.allclose(got, np.asarray([2.0, 3.0, 4.0])) + + +def test_kwonly_defaults() -> None: + expected = X * 2.0 + 0.5 + assert np.allclose(vfn("kwonly_defaults")(X), expected) + expected = X * 4.0 + 1.0 + assert np.allclose(vfn("kwonly_defaults")(X, scale=4.0, bias=1.0), expected) + + +def test_hidden_params_preserve_kwonly_defaults() -> None: + mod = make_module( + "import numpy as np\n" + "C = np.asarray([2.0])\n" + "\n" + "\n" + "def subject(x, *, scale=3):\n" + " return x * scale + C\n" + ) + vec = vectorize(mod.subject) + assert np.allclose(vec(np.asarray([1.0])), [5.0]) + assert np.allclose(vec(np.asarray([1.0]), scale=4), [6.0]) diff --git a/tests/behavior/test_closures_and_helpers.py b/tests/behavior/test_closures_and_helpers.py deleted file mode 100644 index f67ece8..0000000 --- a/tests/behavior/test_closures_and_helpers.py +++ /dev/null @@ -1,259 +0,0 @@ -"""Closures, helper vectorization and caching, name allocation, lambda identification, CSE temps.""" - -from __future__ import annotations - -import math - -import corpus -import numpy as np -import pytest -from support import just_vec, make_fn, make_module, vec_of, vfn - -from array_vectorize import VectorizationError, vectorize - -X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) - - -# ------------------------------------------------------- closures and defaults - - -def test_closure_scalar() -> None: - assert np.allclose(vfn("closure_scalar")(X), X * corpus.SCALE) - - -def test_closure_array() -> None: - got = vfn("closure_array")(np.asarray([1.0, 1.0, 1.0])) - assert np.allclose(got, np.asarray([2.0, 3.0, 4.0])) - - -def test_kwonly_defaults() -> None: - expected = X * 2.0 + 0.5 - assert np.allclose(vfn("kwonly_defaults")(X), expected) - expected = X * 4.0 + 1.0 - assert np.allclose(vfn("kwonly_defaults")(X, scale=4.0, bias=1.0), expected) - - -def test_hidden_params_preserve_kwonly_defaults() -> None: - mod = make_module( - "import numpy as np\n" - "C = np.asarray([2.0])\n" - "\n" - "\n" - "def subject(x, *, scale=3):\n" - " return x * scale + C\n" - ) - vec = vectorize(mod.subject) - assert np.allclose(vec(np.asarray([1.0])), [5.0]) - assert np.allclose(vec(np.asarray([1.0]), scale=4), [6.0]) - - -# ---------------------------------------------------------------- helper calls - - -def test_helper_call_composition() -> None: - mod = make_module( - "import math\n" - "\n" - "\n" - "def square(y):\n" - " return y * y\n" - "\n" - "\n" - "def norm(x):\n" - " return math.sqrt(square(x) + square(x))\n" - ) - vec = vectorize(mod.norm) - xs = np.asarray([3.0, 4.0]) - got = vec(xs) - assert np.allclose(got, np.sqrt(2 * xs**2)) - assert "square_vec(" in vec.source - # the helper itself is vectorized and memoized - from array_vectorize.api import _HELPER_CACHE - - assert (mod.square, False, None) in _HELPER_CACHE - - -def test_helper_partial_application() -> None: - mod = make_module( - "def scale(y, k):\n" - " return y * k\n" - "\n" - "\n" - "def twice(x):\n" - " return scale(x, 2.0) + scale(x, 2.0)\n" - ) - vec = vectorize(mod.twice) - assert np.allclose(vec(np.asarray([1.0, 2.0])), [4.0, 8.0]) - - -def test_helper_with_loop_calling_helper() -> None: - mod = make_module( - "def addmul(a, b):\n" - " s = 0.0\n" - " for i in range(3):\n" - " s = s + a * b\n" - " return s\n" - "\n" - "\n" - "def combined(x, y):\n" - " return addmul(x, y) + addmul(y, x)\n" - ) - vec = vectorize(mod.combined) - xs = np.asarray([2.0]) - ys = np.asarray([3.0]) - assert np.allclose(vec(xs, ys), [2 * 3 * 3 * 2]) - # the loop lives in the helper's generated source, not the caller's - from array_vectorize.api import _HELPER_CACHE - - helper_src = _HELPER_CACHE[(mod.addmul, False, None)].source - assert "for i in range(0, 3):" in helper_src - assert "addmul_vec" in vec.source - - -def test_same_named_helper_does_not_shadow_caller() -> None: - fn = make_fn( - " return helper(x) * 10", - extra=( - "def make_helper():\n" - " def subject(x):\n" - " return x + 1\n" - " return subject\n" - "helper = make_helper()\n" - ), - ) - assert vectorize(fn)(np.asarray([1.0]))[0] == 20.0 - - -def test_helper_cache_protect_flag() -> None: - mod = make_module( - "import math\n\n\n" - "def inner(x):\n" - " return math.sqrt(x) if x > 0 else 0.0\n\n\n" - "def outer(x):\n" - " return inner(x)\n" - ) - vectorize(mod.outer) # caches the unprotected inner - protected = vectorize(mod.outer, protect_domains=True) - with np.errstate(all="raise"): - got = protected(np.asarray([-1.0, 4.0])) - assert np.allclose(got, [0.0, 2.0]) - - -def test_recursive_helper_rejected() -> None: - mod = make_module( - "def fib(n):\n if n < 2:\n return n\n return fib(n - 1) + fib(n - 2)\n" - ) - with pytest.raises(VectorizationError, match="recursive"): - vectorize(mod.fib) - - -def test_helper_keyword_call_rejected() -> None: - fn = make_fn(" return helper(x, y=9)", extra="def helper(x, y=2):\n return x + y\n") - with pytest.raises(VectorizationError, match="keyword"): - vectorize(fn) - - -# ------------------------------------------------------- lambda identification - - -def test_multiple_lambdas_on_one_line() -> None: - mod = make_module("f1, f2 = (lambda x: x + 1), (lambda x: x + 2)\n") - assert vectorize(mod.f2)(np.asarray([1.0]))[0] == 3.0 - assert vectorize(mod.f1)(np.asarray([1.0]))[0] == 2.0 - - -def test_lambdas_differing_only_in_names() -> None: - mod = make_module("import math\nf1, f2 = (lambda x: math.sin(x)), (lambda x: math.cos(x))\n") - assert np.isclose(vectorize(mod.f2)(np.asarray([1.0]))[0], math.cos(1.0)) - assert np.isclose(vectorize(mod.f1)(np.asarray([1.0]))[0], math.sin(1.0)) - - -def test_lambdas_differing_only_in_defaults() -> None: - mod = make_module("g1, g2 = (lambda x, y=1: x + y), (lambda x, y=2: x + y)\n") - assert vectorize(mod.g2)(np.asarray([1.0]))[0] == 3.0 - assert vectorize(mod.g1)(np.asarray([1.0]))[0] == 2.0 - - -def test_lambda_signed_zero_constants() -> None: - mod = make_module( - "import math\n" - "f1, f2 = (\n" - " lambda x: math.copysign(x, 0.0),\n" - " lambda x: math.copysign(x, -0.0),\n" - ")\n" - ) - assert vectorize(mod.f2)(np.asarray([1.0]))[0] == -1.0 - assert vectorize(mod.f1)(np.asarray([1.0]))[0] == 1.0 - - -def test_lambda_int_vs_bool_vs_float_constants() -> None: - mod = make_module("g1, g2, g3 = (lambda x: x + 1), (lambda x: x + True), (lambda x: x + 1.0)\n") - assert vectorize(mod.g2)(np.asarray([1.0]))[0] == 2.0 - assert vectorize(mod.g3)(np.asarray([1.0]))[0] == 2.0 - - -# ------------------------------------------------------------- name allocation - - -def test_reserved_param() -> None: - assert np.allclose(vfn("reserved_param")(X), X + 1) - - -def test_reserved_xp_param_keyword_call() -> None: - mod = make_module("import numpy as np\n\n\ndef subject(xp):\n return xp + 1\n") - vec = vectorize(mod.subject) - assert np.allclose(vec(xp=np.asarray([1.0])), [2.0]) - assert "def subject_vec(xp):" in vec.source - - -def test_namespace_var_rebinding_no_collision() -> None: - mod = make_module( - "import math\n\n\ndef subject(xp):\n xp = xp + 1\n return math.sin(xp)\n" - ) - vec = vectorize(mod.subject) - assert np.allclose(vec(xp=np.asarray([1.0])), [math.sin(2.0)]) - - -def test_reserved_caller_name_param_not_renamed() -> None: - mod = make_module( - "def reserved_caller(reserved_caller_vec):\n return reserved_caller_vec + 1\n" - ) - vec = vectorize(mod.reserved_caller) - assert np.allclose(vec(reserved_caller_vec=np.asarray([1.0])), [2.0]) - - -def test_helper_name_shadowed_by_param() -> None: - vec = vectorize(make_fn(" return min(_vec_common_dtype, 3)", defaults="_vec_common_dtype")) - got = vec(np.asarray([1.0])) - assert got[0] == 1.0 - - -def test_helper_names_shadowed_by_params() -> None: - vec = vectorize( - make_fn( - " return min(_vec_arith_dtype, _vec_common_dtype) + (_vec_arith_dtype > 0)", - defaults="_vec_arith_dtype, _vec_common_dtype=2.0", - ) - ) - got = vec(np.asarray([1.0])) - assert got[0] == 2.0 # min(1, 2) + True - - -# ------------------------------------------------------------------- CSE temps - - -def test_cse_opportunity() -> None: - s = np.sqrt(np.abs(X)) - expected = s * s + s - assert np.allclose(vfn("cse_opportunity")(np.abs(X)), expected) - - -def test_cse_temp_does_not_collide_with_ssa_bindings() -> None: - vec, fn = vec_of(" t = x + 1\n t = t + 1\n return t + math.sin(x) + math.sin(x)") - xs = np.asarray([1.0]) - assert np.isclose(vec(xs)[0], fn(1.0)) - - -def test_signed_zero_not_cse_merged() -> None: - vec = just_vec(" return math.copysign(x, 0.0) + math.copysign(x, -0.0)") - assert np.allclose(vec(np.asarray([1.0])), [0.0]) diff --git a/tests/behavior/test_cse.py b/tests/behavior/test_cse.py new file mode 100644 index 0000000..855a545 --- /dev/null +++ b/tests/behavior/test_cse.py @@ -0,0 +1,28 @@ +"""CSE temps: opportunities, SSA binding collisions, signed zeros.""" + +from __future__ import annotations + +import numpy as np +from support import just_vec, vec_of, vfn + +X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) + + +# ------------------------------------------------------------------- CSE temps + + +def test_cse_opportunity() -> None: + s = np.sqrt(np.abs(X)) + expected = s * s + s + assert np.allclose(vfn("cse_opportunity")(np.abs(X)), expected) + + +def test_cse_temp_does_not_collide_with_ssa_bindings() -> None: + vec, fn = vec_of(" t = x + 1\n t = t + 1\n return t + math.sin(x) + math.sin(x)") + xs = np.asarray([1.0]) + assert np.isclose(vec(xs)[0], fn(1.0)) + + +def test_signed_zero_not_cse_merged() -> None: + vec = just_vec(" return math.copysign(x, 0.0) + math.copysign(x, -0.0)") + assert np.allclose(vec(np.asarray([1.0])), [0.0]) diff --git a/tests/behavior/test_helpers.py b/tests/behavior/test_helpers.py new file mode 100644 index 0000000..3e076a1 --- /dev/null +++ b/tests/behavior/test_helpers.py @@ -0,0 +1,114 @@ +"""Helper vectorization and caching, name shadowing, rejections.""" + +from __future__ import annotations + +import numpy as np +import pytest +from support import make_fn, make_module + +from array_vectorize import VectorizationError, vectorize + +# ---------------------------------------------------------------- helper calls + + +def test_helper_call_composition() -> None: + mod = make_module( + "import math\n" + "\n" + "\n" + "def square(y):\n" + " return y * y\n" + "\n" + "\n" + "def norm(x):\n" + " return math.sqrt(square(x) + square(x))\n" + ) + vec = vectorize(mod.norm) + xs = np.asarray([3.0, 4.0]) + got = vec(xs) + assert np.allclose(got, np.sqrt(2 * xs**2)) + assert "square_vec(" in vec.source + # the helper itself is vectorized and memoized + from array_vectorize.api import _HELPER_CACHE + + assert (mod.square, False, None) in _HELPER_CACHE + + +def test_helper_partial_application() -> None: + mod = make_module( + "def scale(y, k):\n" + " return y * k\n" + "\n" + "\n" + "def twice(x):\n" + " return scale(x, 2.0) + scale(x, 2.0)\n" + ) + vec = vectorize(mod.twice) + assert np.allclose(vec(np.asarray([1.0, 2.0])), [4.0, 8.0]) + + +def test_helper_with_loop_calling_helper() -> None: + mod = make_module( + "def addmul(a, b):\n" + " s = 0.0\n" + " for i in range(3):\n" + " s = s + a * b\n" + " return s\n" + "\n" + "\n" + "def combined(x, y):\n" + " return addmul(x, y) + addmul(y, x)\n" + ) + vec = vectorize(mod.combined) + xs = np.asarray([2.0]) + ys = np.asarray([3.0]) + assert np.allclose(vec(xs, ys), [2 * 3 * 3 * 2]) + # the loop lives in the helper's generated source, not the caller's + from array_vectorize.api import _HELPER_CACHE + + helper_src = _HELPER_CACHE[(mod.addmul, False, None)].source + assert "for i in range(0, 3):" in helper_src + assert "addmul_vec" in vec.source + + +def test_same_named_helper_does_not_shadow_caller() -> None: + fn = make_fn( + " return helper(x) * 10", + extra=( + "def make_helper():\n" + " def subject(x):\n" + " return x + 1\n" + " return subject\n" + "helper = make_helper()\n" + ), + ) + assert vectorize(fn)(np.asarray([1.0]))[0] == 20.0 + + +def test_helper_cache_protect_flag() -> None: + mod = make_module( + "import math\n\n\n" + "def inner(x):\n" + " return math.sqrt(x) if x > 0 else 0.0\n\n\n" + "def outer(x):\n" + " return inner(x)\n" + ) + vectorize(mod.outer) # caches the unprotected inner + protected = vectorize(mod.outer, protect_domains=True) + with np.errstate(all="raise"): + got = protected(np.asarray([-1.0, 4.0])) + assert np.allclose(got, [0.0, 2.0]) + + +def test_recursive_helper_rejected() -> None: + mod = make_module( + "def fib(n):\n if n < 2:\n return n\n return fib(n - 1) + fib(n - 2)\n" + ) + with pytest.raises(VectorizationError, match="recursive"): + vectorize(mod.fib) + + +def test_helper_keyword_call_rejected() -> None: + fn = make_fn(" return helper(x, y=9)", extra="def helper(x, y=2):\n return x + y\n") + with pytest.raises(VectorizationError, match="keyword"): + vectorize(fn) diff --git a/tests/behavior/test_lambdas.py b/tests/behavior/test_lambdas.py index bc2226e..b628fa4 100644 --- a/tests/behavior/test_lambdas.py +++ b/tests/behavior/test_lambdas.py @@ -2,7 +2,10 @@ from __future__ import annotations +import math + import numpy as np +from support import make_module from array_vectorize import vectorize @@ -12,3 +15,42 @@ def test_lambda_vectorize() -> None: got = vec(np.asarray([1.0, 2.0])) assert np.allclose(got, [2.0, 4.0]) assert "def lambda_vec" in vec.source + + +# ------------------------------------------------------- lambda identification + + +def test_multiple_lambdas_on_one_line() -> None: + mod = make_module("f1, f2 = (lambda x: x + 1), (lambda x: x + 2)\n") + assert vectorize(mod.f2)(np.asarray([1.0]))[0] == 3.0 + assert vectorize(mod.f1)(np.asarray([1.0]))[0] == 2.0 + + +def test_lambdas_differing_only_in_names() -> None: + mod = make_module("import math\nf1, f2 = (lambda x: math.sin(x)), (lambda x: math.cos(x))\n") + assert np.isclose(vectorize(mod.f2)(np.asarray([1.0]))[0], math.cos(1.0)) + assert np.isclose(vectorize(mod.f1)(np.asarray([1.0]))[0], math.sin(1.0)) + + +def test_lambdas_differing_only_in_defaults() -> None: + mod = make_module("g1, g2 = (lambda x, y=1: x + y), (lambda x, y=2: x + y)\n") + assert vectorize(mod.g2)(np.asarray([1.0]))[0] == 3.0 + assert vectorize(mod.g1)(np.asarray([1.0]))[0] == 2.0 + + +def test_lambda_signed_zero_constants() -> None: + mod = make_module( + "import math\n" + "f1, f2 = (\n" + " lambda x: math.copysign(x, 0.0),\n" + " lambda x: math.copysign(x, -0.0),\n" + ")\n" + ) + assert vectorize(mod.f2)(np.asarray([1.0]))[0] == -1.0 + assert vectorize(mod.f1)(np.asarray([1.0]))[0] == 1.0 + + +def test_lambda_int_vs_bool_vs_float_constants() -> None: + mod = make_module("g1, g2, g3 = (lambda x: x + 1), (lambda x: x + True), (lambda x: x + 1.0)\n") + assert vectorize(mod.g2)(np.asarray([1.0]))[0] == 2.0 + assert vectorize(mod.g3)(np.asarray([1.0]))[0] == 2.0 diff --git a/tests/behavior/test_name_allocation.py b/tests/behavior/test_name_allocation.py new file mode 100644 index 0000000..2ffef55 --- /dev/null +++ b/tests/behavior/test_name_allocation.py @@ -0,0 +1,59 @@ +"""Generated-name allocation: reserved params, shadowed helper names.""" + +from __future__ import annotations + +import math + +import numpy as np +from support import make_fn, make_module, vfn + +from array_vectorize import vectorize + +X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) + + +# ------------------------------------------------------------- name allocation + + +def test_reserved_param() -> None: + assert np.allclose(vfn("reserved_param")(X), X + 1) + + +def test_reserved_xp_param_keyword_call() -> None: + mod = make_module("import numpy as np\n\n\ndef subject(xp):\n return xp + 1\n") + vec = vectorize(mod.subject) + assert np.allclose(vec(xp=np.asarray([1.0])), [2.0]) + assert "def subject_vec(xp):" in vec.source + + +def test_namespace_var_rebinding_no_collision() -> None: + mod = make_module( + "import math\n\n\ndef subject(xp):\n xp = xp + 1\n return math.sin(xp)\n" + ) + vec = vectorize(mod.subject) + assert np.allclose(vec(xp=np.asarray([1.0])), [math.sin(2.0)]) + + +def test_reserved_caller_name_param_not_renamed() -> None: + mod = make_module( + "def reserved_caller(reserved_caller_vec):\n return reserved_caller_vec + 1\n" + ) + vec = vectorize(mod.reserved_caller) + assert np.allclose(vec(reserved_caller_vec=np.asarray([1.0])), [2.0]) + + +def test_helper_name_shadowed_by_param() -> None: + vec = vectorize(make_fn(" return min(_vec_common_dtype, 3)", defaults="_vec_common_dtype")) + got = vec(np.asarray([1.0])) + assert got[0] == 1.0 + + +def test_helper_names_shadowed_by_params() -> None: + vec = vectorize( + make_fn( + " return min(_vec_arith_dtype, _vec_common_dtype) + (_vec_arith_dtype > 0)", + defaults="_vec_arith_dtype, _vec_common_dtype=2.0", + ) + ) + got = vec(np.asarray([1.0])) + assert got[0] == 2.0 # min(1, 2) + True From ec109de0933fdf76268cd45368447a5b2655d61f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 01:08:47 +0200 Subject: [PATCH 04/17] test: dissolve error paths into per-module suites --- tests/behavior/test_control_flow.py | 13 ++ tests/behavior/test_public_api.py | 15 ++ tests/unit/test_error_paths.py | 211 ---------------------------- tests/unit/test_extract.py | 27 ++++ tests/unit/test_lower.py | 42 ++++++ tests/unit/test_optimize.py | 71 ++++++++++ tests/unit/test_validate.py | 18 +++ 7 files changed, 186 insertions(+), 211 deletions(-) delete mode 100644 tests/unit/test_error_paths.py diff --git a/tests/behavior/test_control_flow.py b/tests/behavior/test_control_flow.py index 1f04c7e..d63d129 100644 --- a/tests/behavior/test_control_flow.py +++ b/tests/behavior/test_control_flow.py @@ -67,3 +67,16 @@ def test_branch_merged_maybe_bool_arithmetic() -> None: ) ) assert list(np.asarray(vec(np.asarray([True, False])))) == [2, 0] + + +# --------------------------------------------------------------- boolop chains + + +def test_deeply_nested_boolop_chain() -> None: + vec = vectorize(make_fn(" return x and y and x and y")) + import numpy as np + + xs = np.asarray([0.0, 1.0, 2.0]) + ys = np.asarray([3.0, 0.0, 5.0]) + out = vec(xs, ys) + assert np.allclose(out, [0.0, 0.0, 5.0]) diff --git a/tests/behavior/test_public_api.py b/tests/behavior/test_public_api.py index 6800762..eb95913 100644 --- a/tests/behavior/test_public_api.py +++ b/tests/behavior/test_public_api.py @@ -2,8 +2,23 @@ from __future__ import annotations +import numpy as np +import pytest + import array_vectorize +from array_vectorize import VectorizationError, vectorize def test_package_exposes_version() -> None: assert array_vectorize.__version__ + + +@pytest.mark.parametrize( + "bad", [[1, 2], {"a": 1}, np.asarray([1.0])], ids=["list", "dict", "array"] +) +def test_unhashable_input_rejected_with_vectorization_error(bad: object) -> None: + # the canonical cache lookup runs before extraction and must not + # surface a TypeError for unhashable junk: the extractor's positioned + # VectorizationError ("expected a plain Python function, got ...") wins + with pytest.raises(VectorizationError, match="plain Python function"): + vectorize(bad) # type: ignore[arg-type] diff --git a/tests/unit/test_error_paths.py b/tests/unit/test_error_paths.py deleted file mode 100644 index 739fd7a..0000000 --- a/tests/unit/test_error_paths.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Coverage-closing tests for error paths and fold branches.""" - -from __future__ import annotations - -import math -from collections.abc import Callable -from typing import Any - -import numpy as np -import pytest -from support import make_module - -from array_vectorize import vectorize -from array_vectorize.errors import VectorizationError -from array_vectorize.frontend.extract import extract_function -from array_vectorize.ir import Binding, BinOp, Call, Literal, Program, Ref, SSAEnv, UnaryOp -from array_vectorize.optimize import const_fold, cse, optimize -from array_vectorize.optimize import dce as dce_pass - - -def make_fn(body: str) -> Callable[..., Any]: - mod = make_module("import math\n\n\ndef subject(x, y=2.0):\n" + body + "\n") - return mod.subject - - -def lit(v: object) -> Literal: - if isinstance(v, bool): - return Literal(v, "bool") - if isinstance(v, int): - return Literal(v, "int") - return Literal(float(v), "float") # type: ignore[arg-type] - - -# ------------------------------------------------------- const-fold branches - - -@pytest.mark.parametrize( - ("op", "a", "b", "expected"), - [ - ("floordiv", 7, 2, 3), - ("mod", 7, 3, 1), - ("and", 6, 3, 2), - ("or", 6, 3, 7), - ("xor", 6, 3, 5), - ("lshift", 1, 3, 8), - ("rshift", 8, 2, 2), - ("sub", 2, 5, -3), - ("mul", 2.5, 4, 10.0), - ], -) -def test_fold_binop_ops(op: str, a: float, b: float, expected: float) -> None: - out = const_fold(Program((), (), BinOp(op, lit(a), lit(b)))) - assert out.result == lit(expected) - - -def test_fold_unary_ops_all() -> None: - assert const_fold(Program((), (), UnaryOp("invert", lit(0)))).result == lit(-1) - assert const_fold(Program((), (), UnaryOp("pos", lit(-2)))).result == lit(-2) - assert const_fold(Program((), (), UnaryOp("neg", lit(2.5)))).result == lit(-2.5) - assert const_fold(Program((), (), UnaryOp("not", lit(True)))).result == lit(False) - - -def test_fold_skips_bad_ops() -> None: - # float bitwise -> TypeError at runtime -> no fold - expr = BinOp("and", lit(1.5), lit(2.5)) - assert const_fold(Program((), (), expr)).result == expr - # unary invert on float -> no fold - expr = UnaryOp("invert", lit(1.5)) - assert const_fold(Program((), (), expr)).result == expr - # unknown unary op passes through - expr = UnaryOp("weird", lit(1)) - assert const_fold(Program((), (), expr)).result == expr - # floordiv by zero - expr = BinOp("floordiv", lit(1), lit(0)) - assert const_fold(Program((), (), expr)).result == expr - - -def test_fold_zero_times_infinity() -> None: - # inf * 0 -> nan: exact IEEE fold - out = const_fold(Program((), (), BinOp("mul", lit(math.inf), lit(0.0)))) - assert math.isnan(out.result.value) # type: ignore[attr-defined] - - -def test_dce_keeps_result_only_binding() -> None: - p = Program(("x",), (Binding("a", Call("abs", (Ref("x"),))),), Ref("a")) - assert dce_pass(p).bindings == p.bindings - - -def test_cse_stops_at_fixpoint() -> None: - inner = Call("sqrt", (Ref("x"),)) - p = Program(("x",), (), BinOp("add", inner, inner)) - out = cse(p, SSAEnv(set())) - assert out.bindings == (Binding("t_1", inner),) - assert out.result == BinOp("add", Ref("t_1"), Ref("t_1")) - - -def test_optimize_no_user_names() -> None: - p = Program(("x",), (Binding("a", Ref("x")),), Ref("a")) - assert optimize(p).result == Ref("a") - - -# ------------------------------------------------------------ extract paths - - -def test_async_function_rejected() -> None: - mod = make_module("async def subject(x):\n return x\n") - with pytest.raises(VectorizationError, match="async"): - extract_function(mod.subject) - - -def test_kwonly_without_default() -> None: - mod = make_module("def subject(x, *, scale):\n return x * scale\n") - info = extract_function(mod.subject) - assert info.params[1].kind == "kwonly" - assert not info.params[1].has_default - - -def test_expr_lambda_top_level() -> None: - make_module("(lambda x: x + 1)\n") - # nothing to grab; just ensure parse path works via direct lambda - info = extract_function(lambda x: x + 1.0) # type: ignore[arg-type] - assert info.params[0].name == "x" - - -def test_no_target_in_source_rejected() -> None: - mod = make_module("X = 1\n") - with pytest.raises(VectorizationError): - extract_function(mod.X) # type: ignore[arg-type] - - -# --------------------------------------------------------- validator paths - - -@pytest.mark.parametrize( - "body", - [ - " y = 1\n return y", - ], -) -def test_validator_ok(body: str) -> None: - extract_function(make_fn(body)) - - -# ------------------------------------------------------------ lower errors - - -def test_calling_math_module_rejected() -> None: - with pytest.raises(VectorizationError, match="math module"): - vectorize(make_fn(" return math(x)")) - - -def test_math_const_called_rejected() -> None: - with pytest.raises(VectorizationError, match="cannot be called"): - vectorize(make_fn(" return math.pi(x)")) - - -def test_calling_shadowed_builtin_rejected() -> None: - with pytest.raises(VectorizationError, match="calling variable"): - vectorize(make_fn(" abs = x\n return abs(x)")) - - -def test_statements_after_return_in_branch_rejected() -> None: - with pytest.raises(VectorizationError, match="after return"): - vectorize(make_fn(" if x > 0:\n return x\n y = 1\n return 0.0")) - - -def test_deeply_nested_boolop_chain() -> None: - vec = vectorize(make_fn(" return x and y and x and y")) - import numpy as np - - xs = np.asarray([0.0, 1.0, 2.0]) - ys = np.asarray([3.0, 0.0, 5.0]) - out = vec(xs, ys) - assert np.allclose(out, [0.0, 0.0, 5.0]) - - -def test_not_all_paths_return_via_branch_fallthrough() -> None: - with pytest.raises(VectorizationError, match="paths return"): - vectorize( - make_fn(" if x > 0:\n return 1.0\n if x < -1:\n return 2.0\n") - ) - - -def test_maybe_unbound_after_nested_branch() -> None: - with pytest.raises(VectorizationError, match="may be unbound"): - vectorize( - make_fn( - " if x > 0:\n" - " if x > 1:\n" - " z = 1.0\n" - " else:\n" - " z = 2.0\n" - " else:\n" - " y = 3.0\n" - " return z" - ) - ) - - -# ------------------------------------------------------ api-level rejections - - -@pytest.mark.parametrize( - "bad", [[1, 2], {"a": 1}, np.asarray([1.0])], ids=["list", "dict", "array"] -) -def test_unhashable_input_rejected_with_vectorization_error(bad: object) -> None: - # the canonical cache lookup runs before extraction and must not - # surface a TypeError for unhashable junk: the extractor's positioned - # VectorizationError ("expected a plain Python function, got ...") wins - with pytest.raises(VectorizationError, match="plain Python function"): - vectorize(bad) # type: ignore[arg-type] diff --git a/tests/unit/test_extract.py b/tests/unit/test_extract.py index baa9107..6809248 100644 --- a/tests/unit/test_extract.py +++ b/tests/unit/test_extract.py @@ -6,6 +6,7 @@ import numpy as np import pytest +from support import make_module from array_vectorize.errors import VectorizationError from array_vectorize.frontend.extract import extract_function @@ -307,3 +308,29 @@ def test_lambda_probe_fallback_rejects_unidentifiable() -> None: tree = ast.parse(inspect.getsource(mod.f1)) with pytest.raises(VectorizationError, match="could not be identified"): _find_target(tree, "", mod.f3, module_source=None) + + +def test_async_function_rejected() -> None: + mod = make_module("async def subject(x):\n return x\n") + with pytest.raises(VectorizationError, match="async"): + extract_function(mod.subject) + + +def test_kwonly_without_default() -> None: + mod = make_module("def subject(x, *, scale):\n return x * scale\n") + info = extract_function(mod.subject) + assert info.params[1].kind == "kwonly" + assert not info.params[1].has_default + + +def test_expr_lambda_top_level() -> None: + make_module("(lambda x: x + 1)\n") + # nothing to grab; just ensure parse path works via direct lambda + info = extract_function(lambda x: x + 1.0) # type: ignore[arg-type] + assert info.params[0].name == "x" + + +def test_no_target_in_source_rejected() -> None: + mod = make_module("X = 1\n") + with pytest.raises(VectorizationError): + extract_function(mod.X) # type: ignore[arg-type] diff --git a/tests/unit/test_lower.py b/tests/unit/test_lower.py index fae25d7..ecb46da 100644 --- a/tests/unit/test_lower.py +++ b/tests/unit/test_lower.py @@ -10,6 +10,7 @@ import pytest from support import make_module +from array_vectorize import vectorize from array_vectorize.errors import VectorizationError from array_vectorize.frontend.extract import extract_function from array_vectorize.ir import ( @@ -515,3 +516,44 @@ def test_kinds_restore_preserves_present_none_kind() -> None: assert k.kind("i") is None k.setdefault_kind("i", "int") # blocked: presence survived the restore assert k.kind("i") is None + + +# ------------------------------------------------------------ lower errors + + +def test_calling_math_module_rejected() -> None: + with pytest.raises(VectorizationError, match="math module"): + vectorize(make_fn(" return math(x)")) + + +def test_math_const_called_rejected() -> None: + with pytest.raises(VectorizationError, match="cannot be called"): + vectorize(make_fn(" return math.pi(x)")) + + +def test_calling_shadowed_builtin_rejected() -> None: + with pytest.raises(VectorizationError, match="calling variable"): + vectorize(make_fn(" abs = x\n return abs(x)")) + + +def test_not_all_paths_return_via_branch_fallthrough() -> None: + with pytest.raises(VectorizationError, match="paths return"): + vectorize( + make_fn(" if x > 0:\n return 1.0\n if x < -1:\n return 2.0\n") + ) + + +def test_maybe_unbound_after_nested_branch() -> None: + with pytest.raises(VectorizationError, match="may be unbound"): + vectorize( + make_fn( + " if x > 0:\n" + " if x > 1:\n" + " z = 1.0\n" + " else:\n" + " z = 2.0\n" + " else:\n" + " y = 3.0\n" + " return z" + ) + ) diff --git a/tests/unit/test_optimize.py b/tests/unit/test_optimize.py index 0427b38..e204574 100644 --- a/tests/unit/test_optimize.py +++ b/tests/unit/test_optimize.py @@ -4,6 +4,8 @@ import math +import pytest + from array_vectorize.ir import ( Binding, BinOp, @@ -17,6 +19,7 @@ Where, ) from array_vectorize.optimize import const_fold, cse, dce, optimize +from array_vectorize.optimize import dce as dce_pass def lit(v: float | int, kind: str = "auto") -> Literal: @@ -216,3 +219,71 @@ def test_optimize_pipeline() -> None: assert out.bindings[0].expr == Call("abs", (folded_sub,)) assert out.bindings[1] == Binding("a", Ref("t_1")) assert out.result == BinOp("add", Ref("t_1"), Ref("a")) + + +# ------------------------------------------------------- const-fold branches + + +@pytest.mark.parametrize( + ("op", "a", "b", "expected"), + [ + ("floordiv", 7, 2, 3), + ("mod", 7, 3, 1), + ("and", 6, 3, 2), + ("or", 6, 3, 7), + ("xor", 6, 3, 5), + ("lshift", 1, 3, 8), + ("rshift", 8, 2, 2), + ("sub", 2, 5, -3), + ("mul", 2.5, 4, 10.0), + ], +) +def test_fold_binop_ops(op: str, a: float, b: float, expected: float) -> None: + out = const_fold(Program((), (), BinOp(op, lit(a), lit(b)))) + assert out.result == lit(expected) + + +def test_fold_unary_ops_all() -> None: + assert const_fold(Program((), (), UnaryOp("invert", lit(0)))).result == lit(-1) + assert const_fold(Program((), (), UnaryOp("pos", lit(-2)))).result == lit(-2) + assert const_fold(Program((), (), UnaryOp("neg", lit(2.5)))).result == lit(-2.5) + assert const_fold(Program((), (), UnaryOp("not", lit(True)))).result == lit(False) + + +def test_fold_skips_bad_ops() -> None: + # float bitwise -> TypeError at runtime -> no fold + expr = BinOp("and", lit(1.5), lit(2.5)) + assert const_fold(Program((), (), expr)).result == expr + # unary invert on float -> no fold + expr = UnaryOp("invert", lit(1.5)) + assert const_fold(Program((), (), expr)).result == expr + # unknown unary op passes through + expr = UnaryOp("weird", lit(1)) + assert const_fold(Program((), (), expr)).result == expr + # floordiv by zero + expr = BinOp("floordiv", lit(1), lit(0)) + assert const_fold(Program((), (), expr)).result == expr + + +def test_fold_zero_times_infinity() -> None: + # inf * 0 -> nan: exact IEEE fold + out = const_fold(Program((), (), BinOp("mul", lit(math.inf), lit(0.0)))) + assert math.isnan(out.result.value) # type: ignore[attr-defined] + + +def test_dce_keeps_result_only_binding() -> None: + p = Program(("x",), (Binding("a", Call("abs", (Ref("x"),))),), Ref("a")) + assert dce_pass(p).bindings == p.bindings + + +def test_cse_stops_at_fixpoint() -> None: + inner = Call("sqrt", (Ref("x"),)) + p = Program(("x",), (), BinOp("add", inner, inner)) + out = cse(p, SSAEnv(set())) + assert out.bindings == (Binding("t_1", inner),) + assert out.result == BinOp("add", Ref("t_1"), Ref("t_1")) + + +def test_optimize_no_user_names() -> None: + p = Program(("x",), (Binding("a", Ref("x")),), Ref("a")) + assert optimize(p).result == Ref("a") diff --git a/tests/unit/test_validate.py b/tests/unit/test_validate.py index fd85310..a9528b5 100644 --- a/tests/unit/test_validate.py +++ b/tests/unit/test_validate.py @@ -341,3 +341,21 @@ def test_unknown_expression_falls_through() -> None: def test_del_is_rejected_by_vectorize() -> None: with pytest.raises(VectorizationError, match="statements are not supported"): vectorize(make_fn(" del x\n return 1.0", signature="x, y=2.0")) + + +# --------------------------------------------------------- validator paths + + +@pytest.mark.parametrize( + "body", + [ + " y = 1\n return y", + ], +) +def test_validator_ok(body: str) -> None: + extract_function(make_fn(body)) + + +def test_statements_after_return_in_branch_rejected() -> None: + with pytest.raises(VectorizationError, match="after return"): + vectorize(make_fn(" if x > 0:\n return x\n y = 1\n return 0.0")) From 9538c0bcd2d22908ecbcea7ae26f7b40592316ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 01:14:35 +0200 Subject: [PATCH 05/17] test: move perf, fuzz and differential suites into their folders --- .agents/skills/add-construct/SKILL.md | 5 +++-- AGENTS.md | 1 + Makefile | 2 +- docs/architecture.md | 15 ++++++++------- docs/development.md | 4 ++-- ...differential.py => test_scalar_equivalence.py} | 0 tests/{ => perf}/test_bench.py | 0 tests/{test_perf.py => perf/test_speed_gate.py} | 0 tests/{test_fuzz_smoke.py => unit/test_fuzz.py} | 0 9 files changed, 15 insertions(+), 12 deletions(-) rename tests/differential/{test_differential.py => test_scalar_equivalence.py} (100%) rename tests/{ => perf}/test_bench.py (100%) rename tests/{test_perf.py => perf/test_speed_gate.py} (100%) rename tests/{test_fuzz_smoke.py => unit/test_fuzz.py} (100%) diff --git a/.agents/skills/add-construct/SKILL.md b/.agents/skills/add-construct/SKILL.md index c50c566..00fd105 100644 --- a/.agents/skills/add-construct/SKILL.md +++ b/.agents/skills/add-construct/SKILL.md @@ -44,8 +44,9 @@ step before starting the next one. and its name to `GOLDEN_NAMES`. Run `uv run pytest --update-golden` and read the new file in `tests/golden/cases/`: it should look like Array API code a person would write. If values are numeric, add a Hypothesis - test in `tests/differential/test_differential.py`. Add unit tests for any - new lowering or optimizer branch so coverage stays at or above 95%. + test in `tests/differential/test_scalar_equivalence.py`. Add unit + tests for any new lowering or optimizer branch so coverage stays at + or above 95%. 6. **Fuzz it.** If the construct is expression-level, add it to the grammar lists in `src/array_vectorize/fuzz.py` (`_UNARY_MATH`, `_BINOPS`, ...). diff --git a/AGENTS.md b/AGENTS.md index b33e80c..71a229f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -57,6 +57,7 @@ the behavior, not the bug that found it). | end-to-end semantics on NumPy | `tests/behavior/test_.py` | | a public option (`fallback`, `protect_domains`, `verify`, `namespace`) | `tests/features/` | | `array-api-strict`, jax, torch | `tests/backends/` | +| benchmarks and the perf gate | `tests/perf/` | | a new corpus function: its generated source | `tests/corpus.py` + `GOLDEN_NAMES`, then `--update-golden` | | a new corpus function: values against scalar Python | `tests/differential/` (Hypothesis) | diff --git a/Makefile b/Makefile index b4ef2ce..56d1346 100644 --- a/Makefile +++ b/Makefile @@ -72,7 +72,7 @@ fuzz: done bench: - $(RUN) pytest tests/test_bench.py -m slow --benchmark-enable \ + $(RUN) pytest tests/perf/test_bench.py -m slow --benchmark-enable \ --benchmark-only --benchmark-columns=min,median,ops \ --benchmark-group-by=func diff --git a/docs/architecture.md b/docs/architecture.md index 47926c8..3853824 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -126,10 +126,12 @@ identity-stable. Tests are organized by level and behavior, never by discovery date: - `tests/unit/` — white-box tests, one file per source module - (frontend, ir, lower, optimize, codegen, runtime). + (frontend, ir, lower, optimize, codegen, runtime), plus the fuzz + smoke test (`test_fuzz.py`). - `tests/behavior/` — black-box end-to-end semantics on NumPy: arithmetic dtype exactness, control flow, loops, calls/math/casts, - closures and helpers, documented divergences. + closures and helpers, inspectability (`test_inspectability.py`: + source/`getsource` attributes), documented divergences. - `tests/features/` — public-option behaviors: fallback, protected domains, differential verification. - `tests/backends/` — array-api-strict compatibility. @@ -138,11 +140,10 @@ Tests are organized by level and behavior, never by discovery date: (`0, ±1, subnormals, ±inf, NaN`). - `tests/golden/` — generated-source snapshots (`cases/`) compared via `ast.dump` equality, immune to formatting drift. -- Top level — `tests/test_inspectability.py` (source/`getsource` - attributes), `tests/test_perf.py` (performance gate: within 5x of - the hand-written array expression and >= 10x faster than - `np.vectorize`), `tests/test_bench.py` (detailed timings, run - `make bench`), and `tests/test_fuzz_smoke.py`. +- `tests/perf/` — the performance gate (`test_speed_gate.py`: within + 5x of the hand-written array expression and >= 10x faster than + `np.vectorize`) and detailed timings (`test_bench.py`, run + `make bench`). The regression corpus from the 26-round independent adversarial review is preserved inside `behavior/` and `features/` (provenance in diff --git a/docs/development.md b/docs/development.md index 4364604..0f80897 100644 --- a/docs/development.md +++ b/docs/development.md @@ -55,8 +55,8 @@ plain `pytest` skips nothing but keeps benchmarks disabled | golden source snapshots | `tests/golden/` (`cases/`, ast.dump equality) | | inspectability (getsource, linecache, memoization, docstrings) | `tests/behavior/test_inspectability.py` etc. | | grammar fuzzer | `src/array_vectorize/fuzz.py` + CI seeds | -| fuzz smoke | `tests/test_fuzz_smoke.py` | -| performance gate + benchmarks | `tests/test_perf.py`, `tests/test_bench.py` | +| fuzz smoke | `tests/unit/test_fuzz.py` | +| performance gate + benchmarks | `tests/perf/` | The historical review-round regression tests are dissolved into `behavior/` and `features/` (provenance in git history). diff --git a/tests/differential/test_differential.py b/tests/differential/test_scalar_equivalence.py similarity index 100% rename from tests/differential/test_differential.py rename to tests/differential/test_scalar_equivalence.py diff --git a/tests/test_bench.py b/tests/perf/test_bench.py similarity index 100% rename from tests/test_bench.py rename to tests/perf/test_bench.py diff --git a/tests/test_perf.py b/tests/perf/test_speed_gate.py similarity index 100% rename from tests/test_perf.py rename to tests/perf/test_speed_gate.py diff --git a/tests/test_fuzz_smoke.py b/tests/unit/test_fuzz.py similarity index 100% rename from tests/test_fuzz_smoke.py rename to tests/unit/test_fuzz.py From 29a85f16480589c05ecc58fd5baf311768cffc02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 01:51:55 +0200 Subject: [PATCH 06/17] test: classify absolute self-imports in the layering test --- tests/unit/test_layering.py | 115 ++++++++++++++++++++++++++++++++---- 1 file changed, 105 insertions(+), 10 deletions(-) diff --git a/tests/unit/test_layering.py b/tests/unit/test_layering.py index 278f87c..d0bfd8c 100644 --- a/tests/unit/test_layering.py +++ b/tests/unit/test_layering.py @@ -1,7 +1,9 @@ """Import-layering rules of the package, checked statically. These encode the boundaries described in docs/architecture.md so a move -that breaks one fails here instead of in review. +that breaks one fails here instead of in review. The classifier is +AST-based, so it does not see dynamic imports (``importlib.import_module``, +``__import__``). """ from __future__ import annotations @@ -14,6 +16,9 @@ PACKAGE = "array_vectorize" ROOT = Path(__file__).parents[2] / "src" / PACKAGE +#: the unit of the package root ``__init__.py`` (its stem; a subpackage's +#: ``__init__.py`` instead belongs to that subpackage's unit) +ROOT_UNIT = "__init__" #: runtime dependencies from pyproject.toml; anything else is dev-only RUNTIME_DEPS = {"array_api_compat"} @@ -27,12 +32,21 @@ def _unit(path: Path) -> str: return parts[0] if len(parts) > 1 else path.stem -def _imports(path: Path) -> tuple[set[str], set[str]]: - """(package units, external top-level modules) imported by ``path``.""" - package_dir = list(path.relative_to(ROOT).parts[:-1]) +def _is_unit(name: str) -> bool: + """Whether ``name`` is a real module or subpackage under the package root.""" + return (ROOT / f"{name}.py").is_file() or (ROOT / name).is_dir() + + +def _classify(tree: ast.Module, package_dir: list[str]) -> tuple[set[str], set[str]]: + """(package units, external top-level modules) imported by ``tree``. + + ``package_dir`` locates the importing file inside the package and + resolves relative imports. Absolute self-imports are policed like + relative ones: ``array_vectorize`` never counts as external. + """ internal: set[str] = set() external: set[str] = set() - for node in ast.walk(ast.parse(path.read_text())): + for node in ast.walk(tree): if isinstance(node, ast.ImportFrom) and node.level: base = package_dir[: len(package_dir) - (node.level - 1)] target = base + (node.module.split(".") if node.module else []) @@ -41,13 +55,38 @@ def _imports(path: Path) -> tuple[set[str], set[str]]: else: internal.update(alias.name for alias in node.names) elif isinstance(node, ast.ImportFrom) and node.module: - external.add(node.module.split(".")[0]) + parts = node.module.split(".") + if parts[0] != PACKAGE: + external.add(parts[0]) + elif len(parts) == 1: + # `from array_vectorize import X, Y`: a real module or + # subpackage X is its own unit; any other name is + # re-exported by the root and reaches the root unit + internal.update( + alias.name if _is_unit(alias.name) else ROOT_UNIT for alias in node.names + ) + else: + internal.add(parts[1]) elif isinstance(node, ast.Import): - external.update(alias.name.split(".")[0] for alias in node.names) + for alias in node.names: + parts = alias.name.split(".") + if parts[0] == PACKAGE: + # importing the package, bare or dotted, reaches at + # least the root unit; the dotted form reaches the + # named unit + internal.add(parts[1] if len(parts) > 1 else ROOT_UNIT) + else: + external.add(parts[0]) external.discard("__future__") return internal, external +def _imports(path: Path) -> tuple[set[str], set[str]]: + """(package units, external top-level modules) imported by ``path``.""" + package_dir = list(path.relative_to(ROOT).parts[:-1]) + return _classify(ast.parse(path.read_text()), package_dir) + + SOURCES = sorted(ROOT.rglob("*.py")) @@ -55,6 +94,55 @@ def _files_of(unit: str) -> list[Path]: return [path for path in SOURCES if _unit(path) == unit] +@pytest.mark.parametrize( + ("statement", "package_dir", "units"), + [ + pytest.param("import array_vectorize", [], {ROOT_UNIT}, id="bare"), + pytest.param("import array_vectorize as av", [], {ROOT_UNIT}, id="bare-aliased"), + pytest.param("from array_vectorize import vectorize", [], {ROOT_UNIT}, id="re-export"), + pytest.param("from array_vectorize import ir", [], {"ir"}, id="submodule"), + pytest.param("from array_vectorize.ir import Node", [], {"ir"}, id="subpackage"), + pytest.param("from array_vectorize.ir.nodes import Node", [], {"ir"}, id="nested"), + pytest.param("import array_vectorize.pipeline", [], {"pipeline"}, id="dotted"), + pytest.param("from . import errors", [], {"errors"}, id="relative-names"), + pytest.param( + "from .errors import VectorizationError", [], {"errors"}, id="relative-module" + ), + pytest.param( + "from .validate import validate", ["frontend"], {"frontend"}, id="relative-inside" + ), + pytest.param( + "from ..errors import VectorizationError", + ["frontend"], + {"errors"}, + id="relative-parent", + ), + ], +) +def test_import_statement_units(statement: str, package_dir: list[str], units: set[str]) -> None: + """Every import form maps to the package units it reaches.""" + internal, external = _classify(ast.parse(statement), package_dir) + assert internal == units + assert PACKAGE not in external + + +@pytest.mark.parametrize("unit", ["runtime", "ir"]) +def test_bare_package_import_breaks_the_leaf_rules(unit: str) -> None: + # `import array_vectorize` classifies as the root unit, an internal + # non-self unit, so the runtime-leaf and ir rules reject it + internal, _ = _classify(ast.parse("import array_vectorize"), [unit]) + assert ROOT_UNIT in internal + assert not internal <= {unit} + + +def test_known_real_imports() -> None: + """The classifier sees the tree's real imports, not only synthetic ones.""" + pipeline_units, _ = _imports(ROOT / "pipeline.py") + assert "lower" in pipeline_units + fuzz_units, _ = _imports(ROOT / "fuzz.py") + assert ROOT_UNIT in fuzz_units + + @pytest.mark.parametrize("path", _files_of("runtime"), ids=lambda p: p.name) def test_runtime_is_a_stdlib_only_leaf(path: Path) -> None: # runtime helpers are injected into generated code and run on the @@ -73,10 +161,17 @@ def test_ir_depends_on_nothing_in_the_package(path: Path) -> None: @pytest.mark.parametrize("path", SOURCES, ids=lambda p: str(p.relative_to(ROOT))) def test_only_the_package_root_imports_api(path: Path) -> None: # the pipeline reaches back into api only through the - # helper_vectorizer callback, never by import + # helper_vectorizer callback, never by import; the root unit + # re-exports api, so a dependency on it reaches api too internal, _ = _imports(path) - if path != ROOT / "__init__.py": - assert "api" not in internal + if path == ROOT / "__init__.py": + return + assert "api" not in internal + if path == ROOT / "fuzz.py": + # the fuzz CLI is a dev tool and deliberately drives the + # compiler through its public API + return + assert ROOT_UNIT not in internal @pytest.mark.parametrize("path", SOURCES, ids=lambda p: str(p.relative_to(ROOT))) From 51de5f407b4ee80bb3b948369ff51b55bb360217 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 01:54:05 +0200 Subject: [PATCH 07/17] style: switch to absolute imports (tid252, ban all) --- pyproject.toml | 8 ++++++++ src/array_vectorize/__init__.py | 4 ++-- src/array_vectorize/api.py | 12 ++++++------ src/array_vectorize/codegen/__init__.py | 2 +- src/array_vectorize/codegen/emitter.py | 8 ++++---- src/array_vectorize/emit.py | 2 +- src/array_vectorize/fallback.py | 4 ++-- src/array_vectorize/frontend/__init__.py | 8 ++++---- src/array_vectorize/frontend/extract.py | 10 +++++----- src/array_vectorize/frontend/lambda_id.py | 4 ++-- src/array_vectorize/frontend/validate.py | 6 +++--- src/array_vectorize/ir/__init__.py | 6 +++--- src/array_vectorize/ir/walk.py | 2 +- src/array_vectorize/lower/__init__.py | 12 ++++++------ src/array_vectorize/lower/env.py | 12 ++++++------ src/array_vectorize/lower/expressions.py | 8 ++++---- src/array_vectorize/lower/helpers.py | 4 ++-- src/array_vectorize/lower/kinds.py | 2 +- src/array_vectorize/lower/loops.py | 6 +++--- src/array_vectorize/lower/sanitize.py | 6 +++--- src/array_vectorize/lower/statements.py | 4 ++-- src/array_vectorize/lower/types.py | 4 ++-- src/array_vectorize/optimize/__init__.py | 12 ++++++------ src/array_vectorize/optimize/constfold.py | 4 ++-- src/array_vectorize/optimize/cse.py | 6 +++--- src/array_vectorize/optimize/dce.py | 4 ++-- src/array_vectorize/optimize/protect_domains.py | 4 ++-- src/array_vectorize/pipeline.py | 14 +++++++------- src/array_vectorize/runtime/__init__.py | 8 ++++---- src/array_vectorize/runtime/arith.py | 2 +- src/array_vectorize/runtime/minmax.py | 7 ++++++- src/array_vectorize/runtime/registry.py | 4 ++-- src/array_vectorize/verify.py | 4 ++-- 33 files changed, 108 insertions(+), 95 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index faff5a3..1420ace 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -96,7 +96,15 @@ select = [ "RUF", # ruff-specific "C4", # flake8-comprehensions "PT", # flake8-pytest-style + "TID252", # flake8-tidy-imports: absolute imports only ] +# the TID252 fix is unsafe-only by default; make `ruff check --fix` apply it +extend-safe-fixes = ["TID252"] + +[tool.ruff.lint.flake8-tidy-imports] +# ruff's default ("parents") only flags `..` imports; every import in the +# tree is absolute +ban-relative-imports = "all" [tool.ruff.lint.per-file-ignores] # scalar functions written for vectorize loop over range(N) for its trip diff --git a/src/array_vectorize/__init__.py b/src/array_vectorize/__init__.py index a4ae8f2..05bc573 100644 --- a/src/array_vectorize/__init__.py +++ b/src/array_vectorize/__init__.py @@ -6,8 +6,8 @@ from __future__ import annotations -from .api import get_source, vectorize -from .errors import VectorizationError +from array_vectorize.api import get_source, vectorize +from array_vectorize.errors import VectorizationError __version__ = "0.1.0" diff --git a/src/array_vectorize/api.py b/src/array_vectorize/api.py index 9ceb801..ea5f4a9 100644 --- a/src/array_vectorize/api.py +++ b/src/array_vectorize/api.py @@ -20,12 +20,12 @@ from collections.abc import Callable from typing import Any -from .compat import _check_namespace -from .errors import VectorizationError -from .fallback import make_fallback -from .frontend.extract import resolve_original -from .pipeline import compile_function -from .verify import verify_match +from array_vectorize.compat import _check_namespace +from array_vectorize.errors import VectorizationError +from array_vectorize.fallback import make_fallback +from array_vectorize.frontend.extract import resolve_original +from array_vectorize.pipeline import compile_function +from array_vectorize.verify import verify_match __all__ = ["get_source", "vectorize"] diff --git a/src/array_vectorize/codegen/__init__.py b/src/array_vectorize/codegen/__init__.py index 0247ed0..c91337a 100644 --- a/src/array_vectorize/codegen/__init__.py +++ b/src/array_vectorize/codegen/__init__.py @@ -2,6 +2,6 @@ from __future__ import annotations -from .emitter import generate_source +from array_vectorize.codegen.emitter import generate_source __all__ = ["generate_source"] diff --git a/src/array_vectorize/codegen/emitter.py b/src/array_vectorize/codegen/emitter.py index 45796df..953f48c 100644 --- a/src/array_vectorize/codegen/emitter.py +++ b/src/array_vectorize/codegen/emitter.py @@ -22,7 +22,8 @@ from inspect import cleandoc from typing import assert_never -from ..ir import ( +from array_vectorize.codegen.docstring import build_docstring, extract_scalar_source +from array_vectorize.ir import ( Binding, BinOp, Call, @@ -38,10 +39,9 @@ Stmt, UnaryOp, Where, - generated_name, # re-export helper + generated_name, ) -from ..lower.types import LoweredFunction -from .docstring import build_docstring, extract_scalar_source +from array_vectorize.lower.types import LoweredFunction __all__ = ["generate_source"] diff --git a/src/array_vectorize/emit.py b/src/array_vectorize/emit.py index 15237d7..66a529e 100644 --- a/src/array_vectorize/emit.py +++ b/src/array_vectorize/emit.py @@ -17,7 +17,7 @@ from collections.abc import Callable from typing import Any, cast -from .ir.ssa import generated_name +from array_vectorize.ir.ssa import generated_name __all__ = ["compile_vectorized"] diff --git a/src/array_vectorize/fallback.py b/src/array_vectorize/fallback.py index d714406..4163df7 100644 --- a/src/array_vectorize/fallback.py +++ b/src/array_vectorize/fallback.py @@ -15,8 +15,8 @@ from array_api_compat import array_namespace -from .codegen.docstring import prefixed_summary -from .compat import _check_namespace, _is_array, _py_scalar +from array_vectorize.codegen.docstring import prefixed_summary +from array_vectorize.compat import _check_namespace, _is_array, _py_scalar __all__ = ["make_fallback"] diff --git a/src/array_vectorize/frontend/__init__.py b/src/array_vectorize/frontend/__init__.py index 41fe9f3..a00608e 100644 --- a/src/array_vectorize/frontend/__init__.py +++ b/src/array_vectorize/frontend/__init__.py @@ -7,10 +7,10 @@ from __future__ import annotations -from .extract import extract_function -from .info import FunctionInfo, Param, ParamKind -from .lambda_id import _find_target -from .tables import ( +from array_vectorize.frontend.extract import extract_function +from array_vectorize.frontend.info import FunctionInfo, Param, ParamKind +from array_vectorize.frontend.lambda_id import _find_target +from array_vectorize.frontend.tables import ( BUILTIN_CASTS, BUILTIN_FOLDS, BUILTIN_UNARY, diff --git a/src/array_vectorize/frontend/extract.py b/src/array_vectorize/frontend/extract.py index 6e76431..85effb7 100644 --- a/src/array_vectorize/frontend/extract.py +++ b/src/array_vectorize/frontend/extract.py @@ -11,11 +11,11 @@ from collections.abc import Callable from typing import Any -from ..compat import _is_array -from ..errors import _reject -from .info import FunctionInfo, Param, ParamKind, _AstFunction -from .lambda_id import _find_target -from .tables import MATH_FUNCS, MATH_SPECIAL +from array_vectorize.compat import _is_array +from array_vectorize.errors import _reject +from array_vectorize.frontend.info import FunctionInfo, Param, ParamKind, _AstFunction +from array_vectorize.frontend.lambda_id import _find_target +from array_vectorize.frontend.tables import MATH_FUNCS, MATH_SPECIAL __all__ = ["FunctionInfo", "Param", "ParamKind", "extract_function", "resolve_original"] diff --git a/src/array_vectorize/frontend/lambda_id.py b/src/array_vectorize/frontend/lambda_id.py index c0d596b..8a5e318 100644 --- a/src/array_vectorize/frontend/lambda_id.py +++ b/src/array_vectorize/frontend/lambda_id.py @@ -14,8 +14,8 @@ from collections.abc import Iterator from typing import Any -from ..errors import _reject -from .info import _AstFunction +from array_vectorize.errors import _reject +from array_vectorize.frontend.info import _AstFunction def _lambdas_in_source_order(tree: ast.AST) -> Iterator[ast.Lambda]: diff --git a/src/array_vectorize/frontend/validate.py b/src/array_vectorize/frontend/validate.py index 7554182..7216e9a 100644 --- a/src/array_vectorize/frontend/validate.py +++ b/src/array_vectorize/frontend/validate.py @@ -9,9 +9,9 @@ import ast -from ..errors import Diagnostic, VectorizationError -from .info import FunctionInfo -from .tables import MATH_CONSTS, MATH_FUNCS, MATH_SPECIAL +from array_vectorize.errors import Diagnostic, VectorizationError +from array_vectorize.frontend.info import FunctionInfo +from array_vectorize.frontend.tables import MATH_CONSTS, MATH_FUNCS, MATH_SPECIAL __all__ = ["validate"] diff --git a/src/array_vectorize/ir/__init__.py b/src/array_vectorize/ir/__init__.py index 59f4650..bd56763 100644 --- a/src/array_vectorize/ir/__init__.py +++ b/src/array_vectorize/ir/__init__.py @@ -2,7 +2,7 @@ from __future__ import annotations -from .nodes import ( +from array_vectorize.ir.nodes import ( Binding, BinOp, Call, @@ -21,8 +21,8 @@ Where, is_bool, ) -from .ssa import RESERVED_NAMES, SSAEnv, generated_name -from .walk import children, rewrite +from array_vectorize.ir.ssa import RESERVED_NAMES, SSAEnv, generated_name +from array_vectorize.ir.walk import children, rewrite __all__ = [ "RESERVED_NAMES", diff --git a/src/array_vectorize/ir/walk.py b/src/array_vectorize/ir/walk.py index 0d6bdbd..099ca29 100644 --- a/src/array_vectorize/ir/walk.py +++ b/src/array_vectorize/ir/walk.py @@ -9,7 +9,7 @@ from collections.abc import Callable from typing import assert_never -from .nodes import ( +from array_vectorize.ir.nodes import ( BinOp, Call, Compare, diff --git a/src/array_vectorize/lower/__init__.py b/src/array_vectorize/lower/__init__.py index 3490d27..c83460b 100644 --- a/src/array_vectorize/lower/__init__.py +++ b/src/array_vectorize/lower/__init__.py @@ -17,12 +17,12 @@ import ast from typing import Any -from ..errors import VectorizationError -from ..frontend.info import FunctionInfo -from ..ir import Node, Program, generated_name -from ..runtime.registry import RUNTIME_HELPERS -from .helpers import _HelperLowerer -from .types import HelperVectorizer, LoweredFunction +from array_vectorize.errors import VectorizationError +from array_vectorize.frontend.info import FunctionInfo +from array_vectorize.ir import Node, Program, generated_name +from array_vectorize.lower.helpers import _HelperLowerer +from array_vectorize.lower.types import HelperVectorizer, LoweredFunction +from array_vectorize.runtime.registry import RUNTIME_HELPERS __all__ = ["HelperVectorizer", "LoweredFunction", "lower_function"] diff --git a/src/array_vectorize/lower/env.py b/src/array_vectorize/lower/env.py index 5248e1b..9606ed2 100644 --- a/src/array_vectorize/lower/env.py +++ b/src/array_vectorize/lower/env.py @@ -14,9 +14,9 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any -from ..errors import VectorizationError -from ..frontend.info import FunctionInfo -from ..ir import ( +from array_vectorize.errors import VectorizationError +from array_vectorize.frontend.info import FunctionInfo +from array_vectorize.ir import ( Binding, Compare, Kind, @@ -28,9 +28,9 @@ Where, is_bool, ) -from ..runtime.registry import RUNTIME_HELPERS -from .kinds import Kinds -from .types import HelperVectorizer +from array_vectorize.lower.kinds import Kinds +from array_vectorize.lower.types import HelperVectorizer +from array_vectorize.runtime.registry import RUNTIME_HELPERS __all__ = ["_LowererBase"] diff --git a/src/array_vectorize/lower/expressions.py b/src/array_vectorize/lower/expressions.py index eb7b66e..1752614 100644 --- a/src/array_vectorize/lower/expressions.py +++ b/src/array_vectorize/lower/expressions.py @@ -12,7 +12,7 @@ import ast -from ..frontend.tables import ( +from array_vectorize.frontend.tables import ( BUILTIN_CASTS, BUILTIN_FOLDS, BUILTIN_UNARY, @@ -20,7 +20,7 @@ MATH_FUNCS, MATH_SPECIAL, ) -from ..ir import ( +from array_vectorize.ir import ( Binding, BinOp, Call, @@ -36,8 +36,8 @@ Where, is_bool, ) -from ..runtime.registry import ArithOp -from .sanitize import _Sanitizer +from array_vectorize.lower.sanitize import _Sanitizer +from array_vectorize.runtime.registry import ArithOp __all__ = ["_ExpressionLowerer"] diff --git a/src/array_vectorize/lower/helpers.py b/src/array_vectorize/lower/helpers.py index 74dc2f4..e27aeec 100644 --- a/src/array_vectorize/lower/helpers.py +++ b/src/array_vectorize/lower/helpers.py @@ -10,8 +10,8 @@ import ast -from ..ir import FuncCall, Node, generated_name -from .loops import _LoopLowerer +from array_vectorize.ir import FuncCall, Node, generated_name +from array_vectorize.lower.loops import _LoopLowerer __all__ = ["_HelperLowerer"] diff --git a/src/array_vectorize/lower/kinds.py b/src/array_vectorize/lower/kinds.py index 1a51a8b..b0e510d 100644 --- a/src/array_vectorize/lower/kinds.py +++ b/src/array_vectorize/lower/kinds.py @@ -45,7 +45,7 @@ from collections.abc import Mapping from dataclasses import dataclass -from ..ir import Kind, Literal +from array_vectorize.ir import Kind, Literal __all__ = ["FactsSnapshot", "Kinds", "VarInfo"] diff --git a/src/array_vectorize/lower/loops.py b/src/array_vectorize/lower/loops.py index 1eab24e..204aeda 100644 --- a/src/array_vectorize/lower/loops.py +++ b/src/array_vectorize/lower/loops.py @@ -12,9 +12,9 @@ import ast -from ..ir import Binding, Kind, Literal, Loop, Node, Ref -from ..optimize.constfold import _const_fold_expr -from .statements import _StatementLowerer +from array_vectorize.ir import Binding, Kind, Literal, Loop, Node, Ref +from array_vectorize.lower.statements import _StatementLowerer +from array_vectorize.optimize.constfold import _const_fold_expr __all__ = ["_LoopLowerer"] diff --git a/src/array_vectorize/lower/sanitize.py b/src/array_vectorize/lower/sanitize.py index 040a040..79186b4 100644 --- a/src/array_vectorize/lower/sanitize.py +++ b/src/array_vectorize/lower/sanitize.py @@ -10,9 +10,9 @@ from __future__ import annotations -from ..ir import Call, DType, Literal, Node, Ref, UnaryOp -from ..ir.walk import children -from .env import _LowererBase +from array_vectorize.ir import Call, DType, Literal, Node, Ref, UnaryOp +from array_vectorize.ir.walk import children +from array_vectorize.lower.env import _LowererBase __all__ = ["_Sanitizer"] diff --git a/src/array_vectorize/lower/statements.py b/src/array_vectorize/lower/statements.py index 61e7104..046505e 100644 --- a/src/array_vectorize/lower/statements.py +++ b/src/array_vectorize/lower/statements.py @@ -10,8 +10,8 @@ import ast -from ..ir import Binding, Literal, Logical, Node, UnaryOp, Where, is_bool -from .expressions import _ExpressionLowerer +from array_vectorize.ir import Binding, Literal, Logical, Node, UnaryOp, Where, is_bool +from array_vectorize.lower.expressions import _ExpressionLowerer __all__ = ["_StatementLowerer"] diff --git a/src/array_vectorize/lower/types.py b/src/array_vectorize/lower/types.py index f03befd..d552160 100644 --- a/src/array_vectorize/lower/types.py +++ b/src/array_vectorize/lower/types.py @@ -12,8 +12,8 @@ from dataclasses import dataclass from typing import Any -from ..frontend.info import Param -from ..ir import Program +from array_vectorize.frontend.info import Param +from array_vectorize.ir import Program __all__ = ["HelperVectorizer", "LoweredFunction"] diff --git a/src/array_vectorize/optimize/__init__.py b/src/array_vectorize/optimize/__init__.py index 8c15399..d126b94 100644 --- a/src/array_vectorize/optimize/__init__.py +++ b/src/array_vectorize/optimize/__init__.py @@ -8,12 +8,12 @@ from __future__ import annotations -from ..ir import Program -from ..ir.ssa import SSAEnv -from .constfold import const_fold -from .cse import cse -from .dce import dce -from .protect_domains import protect_domains +from array_vectorize.ir import Program +from array_vectorize.ir.ssa import SSAEnv +from array_vectorize.optimize.constfold import const_fold +from array_vectorize.optimize.cse import cse +from array_vectorize.optimize.dce import dce +from array_vectorize.optimize.protect_domains import protect_domains __all__ = ["const_fold", "cse", "dce", "optimize", "protect_domains"] diff --git a/src/array_vectorize/optimize/constfold.py b/src/array_vectorize/optimize/constfold.py index 094f4f6..3cfbb4f 100644 --- a/src/array_vectorize/optimize/constfold.py +++ b/src/array_vectorize/optimize/constfold.py @@ -6,7 +6,7 @@ import operator from typing import Any -from ..ir import ( +from array_vectorize.ir import ( Binding, BinOp, Call, @@ -20,7 +20,7 @@ Stmt, UnaryOp, ) -from ..ir.walk import rewrite +from array_vectorize.ir.walk import rewrite __all__ = ["const_fold"] diff --git a/src/array_vectorize/optimize/cse.py b/src/array_vectorize/optimize/cse.py index 809c1bb..42341ce 100644 --- a/src/array_vectorize/optimize/cse.py +++ b/src/array_vectorize/optimize/cse.py @@ -4,9 +4,9 @@ from collections import Counter -from ..ir import Binding, Call, Loop, Node, Program, Ref, Stmt, Where, walk -from ..ir.ssa import SSAEnv -from ..ir.walk import children +from array_vectorize.ir import Binding, Call, Loop, Node, Program, Ref, Stmt, Where, walk +from array_vectorize.ir.ssa import SSAEnv +from array_vectorize.ir.walk import children __all__ = ["cse"] diff --git a/src/array_vectorize/optimize/dce.py b/src/array_vectorize/optimize/dce.py index 926d5cb..9c2e001 100644 --- a/src/array_vectorize/optimize/dce.py +++ b/src/array_vectorize/optimize/dce.py @@ -2,8 +2,8 @@ from __future__ import annotations -from ..ir import Binding, Loop, Node, Program, Ref, Stmt -from ..ir.walk import children +from array_vectorize.ir import Binding, Loop, Node, Program, Ref, Stmt +from array_vectorize.ir.walk import children __all__ = ["dce"] diff --git a/src/array_vectorize/optimize/protect_domains.py b/src/array_vectorize/optimize/protect_domains.py index 94e61ed..e3c01c5 100644 --- a/src/array_vectorize/optimize/protect_domains.py +++ b/src/array_vectorize/optimize/protect_domains.py @@ -2,7 +2,7 @@ from __future__ import annotations -from ..ir import ( +from array_vectorize.ir import ( Binding, Call, Compare, @@ -17,7 +17,7 @@ Where, walk, ) -from ..ir.walk import children +from array_vectorize.ir.walk import children __all__ = ["protect_domains"] diff --git a/src/array_vectorize/pipeline.py b/src/array_vectorize/pipeline.py index fd1c988..b19fec2 100644 --- a/src/array_vectorize/pipeline.py +++ b/src/array_vectorize/pipeline.py @@ -12,13 +12,13 @@ from collections.abc import Callable from typing import Any -from .codegen import generate_source -from .emit import compile_vectorized -from .frontend.extract import extract_function, resolve_original -from .frontend.validate import validate -from .lower import HelperVectorizer, lower_function -from .optimize import optimize, protect_domains -from .verify import verify_match +from array_vectorize.codegen import generate_source +from array_vectorize.emit import compile_vectorized +from array_vectorize.frontend.extract import extract_function, resolve_original +from array_vectorize.frontend.validate import validate +from array_vectorize.lower import HelperVectorizer, lower_function +from array_vectorize.optimize import optimize, protect_domains +from array_vectorize.verify import verify_match __all__ = ["compile_function"] diff --git a/src/array_vectorize/runtime/__init__.py b/src/array_vectorize/runtime/__init__.py index b7f154a..86ce519 100644 --- a/src/array_vectorize/runtime/__init__.py +++ b/src/array_vectorize/runtime/__init__.py @@ -7,9 +7,9 @@ from __future__ import annotations -from .arith import ArithOp, _vec_arith -from .dtype import _fits_dtype -from .minmax import _vec_minmax -from .registry import RUNTIME_HELPERS +from array_vectorize.runtime.arith import ArithOp, _vec_arith +from array_vectorize.runtime.dtype import _fits_dtype +from array_vectorize.runtime.minmax import _vec_minmax +from array_vectorize.runtime.registry import RUNTIME_HELPERS __all__ = ["RUNTIME_HELPERS", "ArithOp", "_fits_dtype", "_vec_arith", "_vec_minmax"] diff --git a/src/array_vectorize/runtime/arith.py b/src/array_vectorize/runtime/arith.py index eb51fcb..317f485 100644 --- a/src/array_vectorize/runtime/arith.py +++ b/src/array_vectorize/runtime/arith.py @@ -10,7 +10,7 @@ import operator from typing import Any -from .dtype import _describe_dtype, _dtype_bits +from array_vectorize.runtime.dtype import _describe_dtype, _dtype_bits __all__ = [ "_ARITH_FNS", diff --git a/src/array_vectorize/runtime/minmax.py b/src/array_vectorize/runtime/minmax.py index 40113fc..54ac89a 100644 --- a/src/array_vectorize/runtime/minmax.py +++ b/src/array_vectorize/runtime/minmax.py @@ -8,7 +8,12 @@ from typing import Any -from .dtype import _describe_dtype, _dtype_bits, _fits_dtype, _minmax_bound_kind +from array_vectorize.runtime.dtype import ( + _describe_dtype, + _dtype_bits, + _fits_dtype, + _minmax_bound_kind, +) __all__ = ["_vec_minmax"] diff --git a/src/array_vectorize/runtime/registry.py b/src/array_vectorize/runtime/registry.py index 15d8c0b..be4e9cf 100644 --- a/src/array_vectorize/runtime/registry.py +++ b/src/array_vectorize/runtime/registry.py @@ -11,8 +11,8 @@ from collections.abc import Callable from typing import Any -from .arith import ArithOp, _vec_arith -from .minmax import _vec_minmax +from array_vectorize.runtime.arith import ArithOp, _vec_arith +from array_vectorize.runtime.minmax import _vec_minmax __all__ = ["RUNTIME_HELPERS", "ArithOp"] diff --git a/src/array_vectorize/verify.py b/src/array_vectorize/verify.py index 92095ad..07a9df3 100644 --- a/src/array_vectorize/verify.py +++ b/src/array_vectorize/verify.py @@ -14,8 +14,8 @@ from array_api_compat import array_namespace -from .compat import _is_array, _py_scalar -from .errors import VectorizationError +from array_vectorize.compat import _is_array, _py_scalar +from array_vectorize.errors import VectorizationError __all__ = ["verify_match"] From d8bf68a425ffc787e0e76f3f2fb57e665dba9605 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 02:12:34 +0200 Subject: [PATCH 08/17] style: adopt Google Python Style (D/N/ICN, google docstrings) --- AGENTS.md | 3 ++ docs/development.md | 19 ++++++++ examples/demo.py | 4 ++ pyproject.toml | 44 ++++++++++++++----- scripts/smoke.py | 3 ++ src/array_vectorize/errors.py | 3 ++ src/array_vectorize/frontend/info.py | 1 + src/array_vectorize/fuzz.py | 2 + src/array_vectorize/ir/nodes.py | 12 +++++ src/array_vectorize/ir/ssa.py | 2 + src/array_vectorize/ir/walk.py | 1 + src/array_vectorize/lower/expressions.py | 25 +++++++---- src/array_vectorize/lower/kinds.py | 20 ++++++++- src/array_vectorize/lower/loops.py | 9 ++-- src/array_vectorize/lower/sanitize.py | 15 ++++--- src/array_vectorize/optimize/constfold.py | 1 + src/array_vectorize/optimize/dce.py | 1 + .../optimize/protect_domains.py | 7 ++- src/array_vectorize/runtime/arith.py | 11 +++-- src/array_vectorize/runtime/dtype.py | 2 +- tests/perf/test_bench.py | 12 ++--- tests/perf/test_speed_gate.py | 8 ++-- tests/unit/test_runtime.py | 5 ++- 23 files changed, 163 insertions(+), 47 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 71a229f..2a22898 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -45,6 +45,9 @@ Invariants: `uv run pytest --update-golden` only when the generated code is meant to change, then review `git diff tests/golden/cases` line by line. +The code follows the Google Python Style Guide, single deviation: +100-column lines. Ruff enforces it via `make lint` / pre-commit. + ## Tests The suite is organized by level and topic. Put each test in the file for diff --git a/docs/development.md b/docs/development.md index 0f80897..3d9b132 100644 --- a/docs/development.md +++ b/docs/development.md @@ -65,6 +65,25 @@ Warnings are errors in the test suite, except numpy's data-dependent numeric warnings (`overflow encountered`, `invalid value`, `divide by zero`) which are expected on documented divergence paths. +## Style + +The code follows the [Google Python Style Guide] +(https://google.github.io/styleguide/pyguide.html). Ruff enforces it in +`make lint`: pydocstyle with the google convention, pep8-naming, import +conventions (`numpy` as `np`), absolute imports only (TID252, +`ban-relative-imports = "all"`), and additionally BLE001, PLW0603, G004, +PGH003 and PGH004 (evaluated: zero violations). + +Evaluated and deliberately skipped: TD (conflicts with Google's current +TODO format) and PLC0415 (import-inside-function is a deliberate pattern +in this compiler's test subjects). + +The only deviation from the guide is the line length: 100 columns (the +guide says 80). Tests relax D1 (missing docstring) — test names are +self-documenting; `tests/corpus.py` ignores all D rules (golden-pinned +test data), and `tests/behavior/test_docstrings.py` ignores D4 (it holds +the NumPy-style docstring subject). + ## CI - **check** — the full `make check` gate on Python 3.12–3.14 diff --git a/examples/demo.py b/examples/demo.py index cf30bf9..2062f8b 100644 --- a/examples/demo.py +++ b/examples/demo.py @@ -9,6 +9,7 @@ @vectorize def sigmoid(x): + """Numerically stable logistic sigmoid.""" if x >= 0: return 1.0 / (1.0 + math.exp(-x)) z = math.exp(x) @@ -17,11 +18,13 @@ def sigmoid(x): @vectorize def gelu(x): + """GELU activation (tanh approximation).""" return 0.5 * x * (1.0 + math.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))) @vectorize def l2_dist(x, y): + """Distance accumulated in a constant-trip loop, then sqrt.""" d = 0.0 for i in range(3): d = d + (x - y) * (x - y) @@ -29,6 +32,7 @@ def l2_dist(x, y): def main() -> None: + """Run the demo on NumPy and array-api-strict inputs.""" x = np.asarray([-3.0, -1.0, 0.0, 1.0, 3.0]) y = np.asarray([2.0, 0.5, -0.5, 0.5, 2.0]) diff --git a/pyproject.toml b/pyproject.toml index 1420ace..dba106c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,30 +86,52 @@ extend-exclude = ["tests/golden/cases", "examples/notebooks"] [tool.ruff.lint] select = [ - "E", # pycodestyle errors - "W", # pycodestyle warnings - "F", # pyflakes - "I", # isort - "UP", # pyupgrade - "B", # flake8-bugbear - "SIM", # flake8-simplify - "RUF", # ruff-specific - "C4", # flake8-comprehensions - "PT", # flake8-pytest-style - "TID252", # flake8-tidy-imports: absolute imports only + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "UP", # pyupgrade + "B", # flake8-bugbear + "SIM", # flake8-simplify + "RUF", # ruff-specific + "C4", # flake8-comprehensions + "PT", # flake8-pytest-style + "TID252", # flake8-tidy-imports: absolute imports only + "D", # pydocstyle (google convention, see below) + "N", # pep8-naming + "ICN", # flake8-import-conventions + # evaluated on the tree: zero violations, kept as guards + "BLE001", # blind except + "PLW0603", # global statement + "G004", # logging f-strings + "PGH003", # blanket type-ignore + "PGH004", # blanket noqa ] # the TID252 fix is unsafe-only by default; make `ruff check --fix` apply it extend-safe-fixes = ["TID252"] +[tool.ruff.lint.pydocstyle] +convention = "google" + [tool.ruff.lint.flake8-tidy-imports] # ruff's default ("parents") only flags `..` imports; every import in the # tree is absolute ban-relative-imports = "all" +[tool.ruff.lint.flake8-import-conventions] +# single source of truth for the import aliases +aliases = { numpy = "np" } + [tool.ruff.lint.per-file-ignores] # scalar functions written for vectorize loop over range(N) for its trip # count; the index is often unused by design "examples/*" = ["B007"] +# test names are self-documenting +"tests/**" = ["D1"] +# corpus functions are golden-pinned test data +"tests/corpus.py" = ["D"] +# holds the NumPy-style docstring subject +"tests/behavior/test_docstrings.py" = ["D4"] [tool.mypy] python_version = "3.12" diff --git a/scripts/smoke.py b/scripts/smoke.py index 9170345..94ee023 100644 --- a/scripts/smoke.py +++ b/scripts/smoke.py @@ -9,6 +9,7 @@ @vectorize def psi(x): + """Exponential decay on the positive half, zero below it.""" if x < 0: return 0.0 return x * math.exp(-x) @@ -16,6 +17,7 @@ def psi(x): @vectorize def clamp(x, lo=0.0, hi=1.0): + """Clamp x into [lo, hi].""" if x < lo: return lo if x > hi: @@ -24,6 +26,7 @@ def clamp(x, lo=0.0, hi=1.0): def main() -> None: + """Run the smoke checks and print the generated sources.""" x = np.asarray([-2.0, -0.5, 0.0, 1.0, 3.0]) expected = np.where(x < 0, 0.0, x * np.exp(-x)) assert np.allclose(psi(x), expected), psi(x) diff --git a/src/array_vectorize/errors.py b/src/array_vectorize/errors.py index 3c776e3..659c35b 100644 --- a/src/array_vectorize/errors.py +++ b/src/array_vectorize/errors.py @@ -16,6 +16,7 @@ class Diagnostic: line: str def render(self) -> str: + """Render the finding as a located snippet with a carets pointer.""" pointer = " " * self.col_offset + "^" return ( f" line {self.lineno}, col {self.col_offset}: {self.message}\n" @@ -32,11 +33,13 @@ class VectorizationError(Exception): """ def __init__(self, message: str, diagnostics: list[Diagnostic] | None = None) -> None: + """Keep the headline message and the collected diagnostics.""" super().__init__(message) self.message = message self.diagnostics: list[Diagnostic] = list(diagnostics or []) def __str__(self) -> str: + """The message, plus every rendered diagnostic below it.""" if not self.diagnostics: return self.message rendered = "\n".join(d.render() for d in self.diagnostics) diff --git a/src/array_vectorize/frontend/info.py b/src/array_vectorize/frontend/info.py index 9feab8d..f0a646a 100644 --- a/src/array_vectorize/frontend/info.py +++ b/src/array_vectorize/frontend/info.py @@ -44,4 +44,5 @@ class FunctionInfo: @property def param_names(self) -> list[str]: + """The parameter names, in signature order.""" return [p.name for p in self.params] diff --git a/src/array_vectorize/fuzz.py b/src/array_vectorize/fuzz.py index 4768ce2..07fe268 100644 --- a/src/array_vectorize/fuzz.py +++ b/src/array_vectorize/fuzz.py @@ -121,6 +121,7 @@ def _statement(rng: random.Random, vars_: list[str], depth: int) -> list[str]: def make_program(rng: random.Random) -> str: + """Build the source of one random scalar function.""" n_vars = rng.randint(1, 2) params = [f"x{i}" for i in range(n_vars)] vars_ = [*params] @@ -201,6 +202,7 @@ def _check_case(rng: random.Random, src: str, fn: Any) -> str: def main(argv: list[str] | None = None) -> int: + """CLI entry point: parse arguments and run the fuzz loop.""" parser = argparse.ArgumentParser(prog="array_vectorize.fuzz", description=__doc__) parser.add_argument("--seconds", type=float, default=60.0, help="time budget") parser.add_argument("--seed", type=int, default=None, help="random seed (reproducible)") diff --git a/src/array_vectorize/ir/nodes.py b/src/array_vectorize/ir/nodes.py index d9de0ae..33294b8 100644 --- a/src/array_vectorize/ir/nodes.py +++ b/src/array_vectorize/ir/nodes.py @@ -58,11 +58,13 @@ class Literal: kind: Kind def __eq__(self, other: object) -> bool: + """Value equality that keeps signed zeros and NaNs distinct.""" if not isinstance(other, Literal): return NotImplemented return self.kind == other.kind and _literal_key(self.value) == _literal_key(other.value) def __hash__(self) -> int: + """Hash consistent with ``__eq__``'s kind/value pair.""" return hash((self.kind, self.value)) @@ -75,6 +77,8 @@ class Ref: @dataclass(frozen=True, slots=True) class BinOp: + """Binary operation, e.g. ``BinOp('add', x, y)`` -> x + y.""" + op: str # add sub mul div pow floordiv mod and or xor lshift rshift left: Node right: Node @@ -82,12 +86,16 @@ class BinOp: @dataclass(frozen=True, slots=True) class UnaryOp: + """Unary operation, e.g. ``UnaryOp('neg', x)`` -> -x.""" + op: str # neg pos invert not operand: Node @dataclass(frozen=True, slots=True) class Compare: + """Comparison, e.g. ``Compare('lt', x, y)`` -> x < y.""" + op: str # eq ne lt le gt ge left: Node right: Node @@ -103,6 +111,8 @@ class Logical: @dataclass(frozen=True, slots=True) class Where: + """Branch merge: ``then`` where ``cond`` holds, ``other`` elsewhere.""" + cond: Node then: Node other: Node @@ -163,6 +173,8 @@ class Loop: @dataclass(frozen=True, slots=True) class Program: + """A lowered function: parameters, top-level bindings, result.""" + params: tuple[str, ...] bindings: tuple[Stmt, ...] result: Node diff --git a/src/array_vectorize/ir/ssa.py b/src/array_vectorize/ir/ssa.py index 50dc859..366c684 100644 --- a/src/array_vectorize/ir/ssa.py +++ b/src/array_vectorize/ir/ssa.py @@ -36,11 +36,13 @@ class SSAEnv: """ def __init__(self, user_names: Iterable[str]) -> None: + """Record the user names the allocator must avoid.""" self._user = frozenset(user_names) self.emitted: set[str] = set() @property def user_names(self) -> frozenset[str]: + """The user names this environment allocates around.""" return self._user def bind(self, var: str, *, force_suffix: bool = False) -> str: diff --git a/src/array_vectorize/ir/walk.py b/src/array_vectorize/ir/walk.py index 099ca29..e5c7447 100644 --- a/src/array_vectorize/ir/walk.py +++ b/src/array_vectorize/ir/walk.py @@ -50,6 +50,7 @@ def rewrite(node: Node, fn: Callable[[Node], Node]) -> Node: def children(node: Node) -> tuple[Node, ...]: + """The direct operand sub-nodes of ``node`` (leaves have none).""" match node: case Literal() | Ref() | DType(): return () diff --git a/src/array_vectorize/lower/expressions.py b/src/array_vectorize/lower/expressions.py index 1752614..2b153d1 100644 --- a/src/array_vectorize/lower/expressions.py +++ b/src/array_vectorize/lower/expressions.py @@ -213,9 +213,11 @@ def _lower_call(self, node: ast.Call) -> Node: raise self.error(node, f"unsupported function {name!r}") def _cast_call(self, arg: Node, dtype: str) -> Node: - """int()/float()/bool()/math.trunc-style casts: fold literal args - (exact), wrap raw loop-variable ints in asarray, pass arrays - through (params are preamble-normalized).""" + """Lower an int()/float()/bool()/math.trunc-style cast. + + Literal args fold exactly, raw loop-variable ints wrap in asarray, + and arrays pass through (params are preamble-normalized). + """ lit = self._literal_value(arg) if lit is not None: try: @@ -254,8 +256,11 @@ def _is_negative_int(node: Node) -> bool: ) def _numeric_kind(self, node: Node) -> Kind | None: - """Best-effort numeric kind: literals, casts, bools, arithmetic, and - Refs whose binding kind was recorded are provable.""" + """Best-effort numeric kind of ``node``. + + Literals, casts, bools, arithmetic, and Refs whose binding kind + was recorded are provable; anything else reads as ``None``. + """ match node: case Literal(): return node.kind @@ -333,11 +338,13 @@ def _numeric_kind(self, node: Node) -> Kind | None: return None def _maybe_bool_result(self, node: Node) -> bool: - """True for min/max results whose static kind is unknown: their - runtime dtype may be boolean (parameter operands), so arithmetic - must use the runtime-polymorphic intify — the _vec_arith_dtype - cast is a no-op for numeric dtypes and int64 for booleans.""" + """True for min/max results whose static kind is unknown. + Their runtime dtype may be boolean (parameter operands), so + arithmetic must use the runtime-polymorphic intify — the + _vec_arith_dtype cast is a no-op for numeric dtypes and int64 + for booleans. + """ if isinstance(node, FuncCall) and node.fn == self.minmax_name: return self._numeric_kind(node) is None if isinstance(node, Ref): diff --git a/src/array_vectorize/lower/kinds.py b/src/array_vectorize/lower/kinds.py index b0e510d..ea78b32 100644 --- a/src/array_vectorize/lower/kinds.py +++ b/src/array_vectorize/lower/kinds.py @@ -75,10 +75,13 @@ class FactsSnapshot: class Kinds: - """Kind/literal/flag facts for every emitted name, plus the stack of - per-loop carried-assignment kind frames.""" + """Kind/literal/flag facts for every emitted name. + + Also holds the stack of per-loop carried-assignment kind frames. + """ def __init__(self) -> None: + """Start with no facts and no active loop frames.""" self._names: dict[str, VarInfo] = {} # names with an established kind fact (a None kind is a legitimate # "unknown" VALUE — presence is what setdefault consults) @@ -101,6 +104,7 @@ def kind(self, name: str) -> Kind | None: return self._names[name].kind def set_kind(self, name: str, kind: Kind | None) -> None: + """Record ``name``'s kind, replacing any previous fact.""" self._kind_established.add(name) self._entry(name).kind = kind @@ -110,19 +114,23 @@ def setdefault_kind(self, name: str, kind: Kind | None) -> None: self.set_kind(name, kind) def update_kinds(self, labels: Mapping[str, Kind | None]) -> None: + """``set_kind`` for every name/kind entry at once.""" for name, kind in labels.items(): self.set_kind(name, kind) # --------------------------------------------------------- literal facts def literal(self, name: str) -> Literal | None: + """The name's recorded literal, or ``None``.""" info = self._names.get(name) return info.literal if info is not None else None def set_literal(self, name: str, lit: Literal) -> None: + """Record ``name`` as holding the literal ``lit``.""" self._entry(name).literal = lit def drop_literal(self, name: str) -> None: + """Forget ``name``'s literal fact, if any.""" info = self._names.get(name) if info is not None: info.literal = None @@ -130,25 +138,31 @@ def drop_literal(self, name: str) -> None: # ----------------------------------------------------------------- flags def is_scalar(self, name: str) -> bool: + """True when the name may hold a raw Python scalar.""" info = self._names.get(name) return info is not None and (info.maybe_scalar or info.literal is not None) def mark_scalar(self, name: str) -> None: + """Mark ``name`` as possibly holding a raw Python scalar.""" self._entry(name).maybe_scalar = True def unmark_scalar(self, name: str) -> None: + """Clear the maybe-scalar mark on ``name``.""" info = self._names.get(name) if info is not None: info.maybe_scalar = False def is_maybe_bool(self, name: str) -> bool: + """True when the name may hold a boolean.""" info = self._names.get(name) return info is not None and info.maybe_bool def mark_maybe_bool(self, name: str) -> None: + """Mark ``name`` as possibly boolean.""" self._entry(name).maybe_bool = True def unmark_maybe_bool(self, name: str) -> None: + """Clear the maybe-boolean mark on ``name``.""" info = self._names.get(name) if info is not None: info.maybe_bool = False @@ -156,9 +170,11 @@ def unmark_maybe_bool(self, name: str) -> None: # -------------------------------------------------------- carried frames def push_carried_frame(self) -> None: + """Enter a loop: body assignments record into a fresh frame.""" self.carried_frames.append({}) def pop_carried_frame(self) -> dict[str, set[Kind | None]]: + """Leave a loop, returning its carried-assignment frame.""" return self.carried_frames.pop() def record_carried_kind(self, name: str, kind: Kind | None) -> None: diff --git a/src/array_vectorize/lower/loops.py b/src/array_vectorize/lower/loops.py index 204aeda..ea0aa8b 100644 --- a/src/array_vectorize/lower/loops.py +++ b/src/array_vectorize/lower/loops.py @@ -133,9 +133,12 @@ def lower_for(self, stmt: ast.For) -> None: def union_labels( frame: dict[str, set[Kind | None]], ) -> tuple[dict[str, Kind | None], list[str]]: - """Per carried name: the final kind (union of the phi kind and - all body-assignment kinds) and the names whose kinds MIX across - iterations.""" + """Per carried name: the final kind and the mixed names. + + The final kind is the union of the phi kind and all + body-assignment kinds; a name is mixed when its kinds MIX + across iterations. + """ labels: dict[str, Kind | None] = {} mixed: list[str] = [] for loop_carried_name in carried.values(): diff --git a/src/array_vectorize/lower/sanitize.py b/src/array_vectorize/lower/sanitize.py index 79186b4..c5bd3e1 100644 --- a/src/array_vectorize/lower/sanitize.py +++ b/src/array_vectorize/lower/sanitize.py @@ -19,8 +19,11 @@ class _Sanitizer(_LowererBase): def _literal_value(self, arg: Node) -> Literal | None: - """The literal behind an arg: directly, as a negated literal - (``-1`` lowers to UnaryOp(neg, 1)), or via a literal binding.""" + """The literal behind an arg, or ``None``. + + Found directly, as a negated literal (``-1`` lowers to + UnaryOp(neg, 1)), or via a literal binding. + """ if isinstance(arg, Literal): return arg if ( @@ -35,11 +38,13 @@ def _literal_value(self, arg: Node) -> Literal | None: return None def _possibly_scalar(self, node: Node) -> bool: - """True when the node's runtime value may be a raw Python scalar: - it is a literal, a possibly-scalar name (parameter default, loop + """True when the node's runtime value may be a raw Python scalar. + + It is a literal, a possibly-scalar name (parameter default, loop variable, literal binding), or an expression computed from one. Operators promote such scalars fine; xp.* function arguments do - not (strict backends reject plain scalars).""" + not (strict backends reject plain scalars). + """ def walk(n: Node) -> bool: if isinstance(n, Literal): diff --git a/src/array_vectorize/optimize/constfold.py b/src/array_vectorize/optimize/constfold.py index 3cfbb4f..0ce5cbe 100644 --- a/src/array_vectorize/optimize/constfold.py +++ b/src/array_vectorize/optimize/constfold.py @@ -241,5 +241,6 @@ def _fold_stmt(stmt: Stmt) -> Stmt: def const_fold(program: Program) -> Program: + """Constant-fold every binding and the result expression.""" bindings = tuple(_fold_stmt(s) for s in program.bindings) return Program(program.params, bindings, _const_fold_expr(program.result)) diff --git a/src/array_vectorize/optimize/dce.py b/src/array_vectorize/optimize/dce.py index 9c2e001..ba731e9 100644 --- a/src/array_vectorize/optimize/dce.py +++ b/src/array_vectorize/optimize/dce.py @@ -50,6 +50,7 @@ def _mark_live_stmt(stmt: Stmt, live: set[str]) -> Stmt | None: def dce(program: Program) -> Program: + """Drop bindings whose results are never used.""" live: set[str] = set() _mark_live(program.result, live) kept: list[Stmt] = [] diff --git a/src/array_vectorize/optimize/protect_domains.py b/src/array_vectorize/optimize/protect_domains.py index e3c01c5..4574375 100644 --- a/src/array_vectorize/optimize/protect_domains.py +++ b/src/array_vectorize/optimize/protect_domains.py @@ -148,8 +148,11 @@ def mark_fully_live(node: Node) -> None: mark_fully_live(child) def rewrite_stmt(stmt: Stmt, ahead: set[str]) -> Stmt: - """Rewrite one binding; ``ahead`` holds the names bound AFTER this - statement (they are not yet assigned where this statement runs).""" + """Rewrite one binding. + + ``ahead`` holds the names bound AFTER this statement (they are + not yet assigned where this statement runs). + """ if isinstance(stmt, Binding): ctx = live.pop(stmt.name, None) record_ctx = ctx diff --git a/src/array_vectorize/runtime/arith.py b/src/array_vectorize/runtime/arith.py index 317f485..4971144 100644 --- a/src/array_vectorize/runtime/arith.py +++ b/src/array_vectorize/runtime/arith.py @@ -61,8 +61,10 @@ def _is_u64(a: Any) -> bool: def _to_u64(xp: Any, a: Any) -> Any: - """Cast to uint64, wrapping negative int literals modularly (NumPy - rejects out-of-bounds Python ints in asarray).""" + """Cast to uint64, wrapping negative int literals modularly. + + NumPy rejects out-of-bounds Python ints in asarray. + """ if isinstance(a, int) and not isinstance(a, bool) and a < 0: a = a % 2**64 return xp.asarray(a, dtype=xp.uint64) @@ -105,8 +107,9 @@ def _u64_sub_exact(xp: Any, left: Any, right: Any) -> Any | None: def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any: - """Exact-semantics arithmetic for bool-intified operations (injected - into generated modules). ``op_id`` indexes _ARITH_OPS. + """Exact-semantics arithmetic for bool-intified operations. + + Injected into generated modules; ``op_id`` indexes _ARITH_OPS. Bools are cast to a numeric dtype (int64 headroom; Python integers are unbounded). Floats keep their dtype; mixed int/float promotes to diff --git a/src/array_vectorize/runtime/dtype.py b/src/array_vectorize/runtime/dtype.py index 5a424d7..fbe4a51 100644 --- a/src/array_vectorize/runtime/dtype.py +++ b/src/array_vectorize/runtime/dtype.py @@ -82,7 +82,7 @@ def _dtype_bits(dtype: Any) -> int: def _minmax_bound_kind(lit: Any, is_min: Any, dtype: Any) -> str: - """ "fit", "clamp", or "promote" for a literal min/max bound. + """``"fit"``, ``"clamp"``, or ``"promote"`` for a literal min/max bound. A bound that cannot win is CLAMPED into an unsigned array dtype instead of forcing promotion: max(unsigned, negative) never selects diff --git a/tests/perf/test_bench.py b/tests/perf/test_bench.py index 03a1826..58e04b5 100644 --- a/tests/perf/test_bench.py +++ b/tests/perf/test_bench.py @@ -19,7 +19,7 @@ import math -import corpus as CORPUS +import corpus import numpy as np import pytest @@ -32,11 +32,11 @@ # representative kernels: plain arithmetic, branch-heavy, loop-accumulated # (name, function, positional argument count) KERNELS = [ - ("add", CORPUS.add, 2), - ("arith_ops", CORPUS.arith_ops, 2), - ("relu", CORPUS.relu, 1), - ("psi", CORPUS.psi, 1), - ("clamp", CORPUS.clamp, 1), + ("add", corpus.add, 2), + ("arith_ops", corpus.arith_ops, 2), + ("relu", corpus.relu, 1), + ("psi", corpus.psi, 1), + ("clamp", corpus.clamp, 1), ] SIZES = [1_000, 100_000, 1_000_000] diff --git a/tests/perf/test_speed_gate.py b/tests/perf/test_speed_gate.py index 60ac7d5..8a4afc3 100644 --- a/tests/perf/test_speed_gate.py +++ b/tests/perf/test_speed_gate.py @@ -27,9 +27,11 @@ def _best_of(fn: Callable[[], object], repeats: int = 3) -> float: - """Best (minimum) wall time over repeats - the standard noise-robust - timing statistic: scheduler hiccups and cache pollution only ever make - a run slower.""" + """Best (minimum) wall time over ``repeats`` runs. + + The standard noise-robust timing statistic: scheduler hiccups and + cache pollution only ever make a run slower. + """ best = math.inf for _ in range(repeats): t0 = time.perf_counter() diff --git a/tests/unit/test_runtime.py b/tests/unit/test_runtime.py index e31309a..cd86e45 100644 --- a/tests/unit/test_runtime.py +++ b/tests/unit/test_runtime.py @@ -9,12 +9,15 @@ from __future__ import annotations import numpy as np -import numpy as xp from array_vectorize.runtime.arith import ArithOp, _vec_arith from array_vectorize.runtime.dtype import _describe_dtype, _fits_dtype from array_vectorize.runtime.minmax import _vec_minmax +# the array namespace passed to the runtime helpers, as generated code +# binds it +xp = np + U64_MAX = 2**64 - 1 I64_MAX = 2**63 - 1 From 9ea5333bca3061088203e9e77659244d4b25571d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 03:47:24 +0200 Subject: [PATCH 09/17] build: adopt cleanporter for Google-style imports Wire cleanporter 0.6.0 into the locked environment, pyproject config, pre-commit and `make fmt`; the import rewrite itself follows in the next commits. - pinned to the local stand-in git source at e53776c, the commit carrying the PEP 695 type-alias fix; the upstream 0.6.0 release is pending, so the source is swapped to the GitHub URL + tag once it is cut - the lint gate line and treat_unresolved_as_error land together with the import rewrite in the following commits, so every intermediate commit stays green; check mode currently reports 385 CP001 / 4 CP002 / 2 CP003 findings - the fmt recipe runs the fixers in the exit-code-robust order: ruff format first and last, ruff check --fix --exit-zero in between, and cleanporter --fix under `-` because it exits 1 whenever anything remains it will not rewrite - golden cases, marimo notebooks and tests/corpus.py are exempt: pinned compiler output, notebook format owned by marimo, and golden test data whose from-imports deliberately exercise the supported subset --- .pre-commit-config.yaml | 6 +++ Makefile | 10 ++++- pyproject.toml | 13 ++++++ uv.lock | 88 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 116 insertions(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index b5753f4..a577b82 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,6 +18,12 @@ repos: # never disagree on the ruff version - repo: local hooks: + - id: cleanporter-fix + name: cleanporter fix + entry: uv run --frozen cleanporter --fix --whole-project + language: system + types_or: [python, pyi] + require_serial: true - id: ruff-check name: ruff check entry: uv run --frozen ruff check --fix --force-exclude diff --git a/Makefile b/Makefile index 56d1346..20ddd68 100644 --- a/Makefile +++ b/Makefile @@ -18,8 +18,16 @@ install: $(UV) sync $(RUN) pre-commit install +# fixers run to completion whatever findings remain; `make lint` stays +# the gate. `-` on cleanporter: it exits 1 whenever anything remains it +# will not rewrite (no exit-zero flag). ruff check --fix gets +# --exit-zero: ruff cannot fix E501 itself, and cleanporter's module +# prefixes push lines past 100 columns; ruff format runs first and last +# so the fixes land on wrapped code fmt: - $(RUN) ruff check --fix $(PY_PATHS) + -$(RUN) cleanporter --fix $(PY_PATHS) + $(RUN) ruff format $(PY_PATHS) + $(RUN) ruff check --fix --exit-zero $(PY_PATHS) $(RUN) ruff format $(PY_PATHS) lint: diff --git a/pyproject.toml b/pyproject.toml index dba106c..2c9c1aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,7 @@ dev = [ "numpy>=2.0", "array-api-strict>=2.3", "pre-commit>=3.5", + "cleanporter>=0.6", ] backends = [ "jax>=0.4", @@ -71,6 +72,9 @@ explicit = true [tool.uv.sources] torch = { index = "pytorch-cpu" } +# local stand-in pin for the PEP 695 type-alias fix; swap to the GitHub +# URL + 0.6.0 tag when the upstream release is cut +cleanporter = { git = "file:///tmp/opencode/cleanporter", rev = "e53776c" } [tool.hatch.build.targets.wheel] packages = ["src/array_vectorize"] @@ -133,6 +137,15 @@ aliases = { numpy = "np" } # holds the NumPy-style docstring subject "tests/behavior/test_docstrings.py" = ["D4"] +[tool.cleanporter] +scope = "all" +# golden cases are pinned compiler output; marimo owns the notebook format +exclude = ["tests/golden/cases/", "examples/notebooks/"] + +[[tool.cleanporter.skip]] +file = 'tests/corpus\.py' +reason = "golden test data: from-imports deliberately exercise the supported subset" + [tool.mypy] python_version = "3.12" strict = true diff --git a/uv.lock b/uv.lock index f69f3cd..4cc61d2 100644 --- a/uv.lock +++ b/uv.lock @@ -62,6 +62,7 @@ backends = [ ] dev = [ { name = "array-api-strict" }, + { name = "cleanporter" }, { name = "coverage" }, { name = "hypothesis" }, { name = "mypy" }, @@ -91,6 +92,7 @@ backends = [ ] dev = [ { name = "array-api-strict", specifier = ">=2.3" }, + { name = "cleanporter", git = "file:///tmp/opencode/cleanporter?rev=e53776c" }, { name = "coverage", extras = ["toml"], specifier = ">=7.6" }, { name = "hypothesis", specifier = ">=6.100" }, { name = "mypy", specifier = ">=1.20" }, @@ -266,6 +268,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/db/3c/33bac158f8ab7f89b2e59426d5fe2e4f63f7ed25df84c036890172b412b5/cfgv-3.5.0-py2.py3-none-any.whl", hash = "sha256:a8dc6b26ad22ff227d2634a65cb388215ce6cc96bbcc5cfde7641ae87e8dacc0", size = 7445, upload-time = "2025-11-19T20:55:50.744Z" }, ] +[[package]] +name = "cleanporter" +version = "0.6.0" +source = { git = "file:///tmp/opencode/cleanporter?rev=e53776c#e53776c0898cf3fe0ac141fa2ce3ff4655213520" } +dependencies = [ + { name = "libcst" }, +] + [[package]] name = "click" version = "8.5.0" @@ -623,6 +633,60 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, ] +[[package]] +name = "libcst" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml", marker = "python_full_version != '3.13.*'" }, + { name = "pyyaml-ft", marker = "python_full_version == '3.13.*'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/02/c0/098e5c91ff1537f00c85a6438b6cb1863d17144680cc91f47c87f104a200/libcst-1.9.0.tar.gz", hash = "sha256:087b58a9afe076bb08e2d726478e1f16cb928d67ffa9092817e033c335de522a", size = 914739, upload-time = "2026-07-29T21:28:43.153Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b0/bb/d22c37c33dfe18084634f5ef89f8f0749ffe7b6e0ad312722aafd86bbbdb/libcst-1.9.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:cd1a3500c41784075c4946a995d5ad89f68fa0d226b63ff3c4d78f6ea6dd23e5", size = 2043159, upload-time = "2026-07-29T19:24:49.013Z" }, + { url = "https://files.pythonhosted.org/packages/10/b8/2dedef84d72e7271119217503b69ed6dc5d0b2077685e163caae669d9c70/libcst-1.9.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:611cebd3bbc2014576f4dcc7b845b3c594c96ddc287a3db9b78f22eff156a7d3", size = 2203245, upload-time = "2026-07-29T19:24:50.399Z" }, + { url = "https://files.pythonhosted.org/packages/e8/90/e02ac2dad647423f947bb11f8322bfdba8ccfdd380e6c7b695add2d1acd4/libcst-1.9.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:8d731abe1307720ea1a52d447555e8443a6d130e0e520243c0634a58f6edbc9d", size = 2255388, upload-time = "2026-07-29T19:24:52.21Z" }, + { url = "https://files.pythonhosted.org/packages/13/5f/6089a51518cfd2ff40950eb26bcc36951d7b6d4f4213568aa0290265aff7/libcst-1.9.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8bc5351d92ca6ac1cc32e097700e1161ad1ceaa4d9b2cca5abadb1e94576b325", size = 2268982, upload-time = "2026-07-29T19:24:53.562Z" }, + { url = "https://files.pythonhosted.org/packages/ed/78/26881ec466fb70cbc129dca26ccb5a52a0061face2c5822e4b61f00f9699/libcst-1.9.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:03165a264653bb77f6a11b412ae09c08bdb0c25864f3b8d42b816ac64b9d4b9e", size = 2378174, upload-time = "2026-07-29T19:24:55.175Z" }, + { url = "https://files.pythonhosted.org/packages/e1/7a/a4dba5f11faf12a12ffba06d18019987851aab33b590242a602ffd4fb1bd/libcst-1.9.0-cp312-cp312-win_amd64.whl", hash = "sha256:b755ed4a4bc2faee849b54820023137d60d4c199e0a0822f0ff0c1bc49e49b48", size = 2103462, upload-time = "2026-07-29T19:24:56.966Z" }, + { url = "https://files.pythonhosted.org/packages/f5/13/57cb129093e0d6744b3c914ba6cbbdf51ed1b2c823882936c553a3a0cf1b/libcst-1.9.0-cp312-cp312-win_arm64.whl", hash = "sha256:6e50576bad7d56459d9792cd0b0dfe5469dad13646e9a9c0a8b1ba20b269f332", size = 1979825, upload-time = "2026-07-29T19:24:58.403Z" }, + { url = "https://files.pythonhosted.org/packages/f2/b1/befc0544283bb3923a928accf79ed685e5a725524bdb3491826670affc07/libcst-1.9.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:b8b9df30f524317b097dc53065b25dda33d6a4cc3c7c8bf4fc83ca7559c58cc0", size = 2043754, upload-time = "2026-07-29T19:24:59.885Z" }, + { url = "https://files.pythonhosted.org/packages/45/50/fef7c172a8457c95894edf5fb04805024899cdfea41fa01b0636587b79e1/libcst-1.9.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:e465a7bc9c2b9533eb9e06d2391f8819f811b5112919d4064c9fb8565aaafa08", size = 2203548, upload-time = "2026-07-29T19:25:01.323Z" }, + { url = "https://files.pythonhosted.org/packages/18/ff/764cd2be1fd99d774fc44039c319dc0ed1d9d9afeaa02759a05121cccd4b/libcst-1.9.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:8504b422c95676a8c27b517e1ac01413ece91bf356865c587ca9bdcd5708a2f7", size = 2255274, upload-time = "2026-07-29T19:25:02.764Z" }, + { url = "https://files.pythonhosted.org/packages/34/a7/474748a27a02fa83e3556b260d5f5236ca48164fa7beffecf3b2cad24dca/libcst-1.9.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:bcb9f9d4fcfe2ec7a40d2c26e03a538d5d9dc189c38551eb3f94ab661afee7c0", size = 2269126, upload-time = "2026-07-29T19:25:04.158Z" }, + { url = "https://files.pythonhosted.org/packages/8a/b7/655e45363b8cf87b91e41e060b21c89e5316c7ef36eb14a1a498b27bb71e/libcst-1.9.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e8d671c39a431c309476099b8ec811e412503ec0f4465f6fc907cb51c70e8e6b", size = 2378602, upload-time = "2026-07-29T19:25:06.424Z" }, + { url = "https://files.pythonhosted.org/packages/db/13/6da63f0902ece43bf9d737017251ad7ad06edd9d3c4e1856450403cb4473/libcst-1.9.0-cp313-cp313-win_amd64.whl", hash = "sha256:4d382fba04077eb556a1ea4295a4482e192aa93639c5a07ba0885841965ea0c0", size = 2103486, upload-time = "2026-07-29T19:25:07.979Z" }, + { url = "https://files.pythonhosted.org/packages/fd/ff/dccb1a55e38b4e47256a97242d61d2bef5366c744faf88fb087d4fa0f995/libcst-1.9.0-cp313-cp313-win_arm64.whl", hash = "sha256:a621e261990148c1cfbe26c1798bd2375c131cf358a44a7a4659fc41c1a336e1", size = 1979907, upload-time = "2026-07-29T19:25:09.532Z" }, + { url = "https://files.pythonhosted.org/packages/65/2a/4943c71d90975bc59034a057dea346b365c276f308c1d31f5cf2bd85492d/libcst-1.9.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:eccf4c57d273cdd3fe1c67b72cf9bb1bbd4547aa011824e96ccf5b7136057aa4", size = 2044529, upload-time = "2026-07-29T19:25:11.04Z" }, + { url = "https://files.pythonhosted.org/packages/be/ed/1168b98c2a0f338be3a44753647baab051feaf6167cc887bf4447f8fd920/libcst-1.9.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:32395244edfe6538e0ea2bf82051d60103d3f54861274805c4fb3745efb70a85", size = 2203519, upload-time = "2026-07-29T19:25:12.543Z" }, + { url = "https://files.pythonhosted.org/packages/25/7d/2eaa697a80f899bcf2245680bbfcc1e478eb3b3328c21d4b2880bdf00dac/libcst-1.9.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:444e84c76cd035cd2fe136838c1a524d34b08216521f6f5093df8a6f6cfa5799", size = 2255879, upload-time = "2026-07-29T19:25:14.137Z" }, + { url = "https://files.pythonhosted.org/packages/90/03/793b9fd96dd52d202f5f2b88d7e54f05ebed767d1bb3b32395a1817bf6ab/libcst-1.9.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:bb5d0946f2b4c6711b5d69fe4f833b364f9e7a1a2b08f88b98619dc18975099a", size = 2269807, upload-time = "2026-07-29T19:25:15.647Z" }, + { url = "https://files.pythonhosted.org/packages/a3/f4/1bc7aaea03971c45e8a885fed9fc1c73176dbbd00b0296a3cde9529bafd6/libcst-1.9.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:45808c03528b3ad40b14095348a918e08d98c47b4a125d631cb78e817c0a5b16", size = 2378171, upload-time = "2026-07-29T19:25:17.289Z" }, + { url = "https://files.pythonhosted.org/packages/8e/f7/bc49e367d2bc8817213594dd52cf6b2da2dbbda90b5382d60653999274ac/libcst-1.9.0-cp314-cp314-win_amd64.whl", hash = "sha256:568288cdbfe3b4ca3ae4852cb0a439ff053dd54c841bc4995bf2b71238b5de40", size = 2177847, upload-time = "2026-07-29T19:25:19.35Z" }, + { url = "https://files.pythonhosted.org/packages/87/80/4d81577a22e6d535d1a3409f3a1c6903e09036f0e17fc47f92169dbc3501/libcst-1.9.0-cp314-cp314-win_arm64.whl", hash = "sha256:107593af46945593e7825821793393262bc4fa1d3ea24c3ed487b61269bbbdf8", size = 2058134, upload-time = "2026-07-29T19:25:20.769Z" }, + { url = "https://files.pythonhosted.org/packages/d1/7f/c3f3a0e7a1a2adaa76e815e82a7814d6bfe7d49432276bba652248c68d0b/libcst-1.9.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f6248cb07444ab9a6733a855737a9febed8b9adca51347019348b09a3ac7dfe9", size = 2036102, upload-time = "2026-07-29T19:25:22.212Z" }, + { url = "https://files.pythonhosted.org/packages/4e/af/2f5543255b2c7d749b966adeccc6bfb1652cf8b425afb913d0f3f025a865/libcst-1.9.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:496c24e0d3240bc7da45dae543aff3f5b6509978c39262d6b84ed2fb999dded2", size = 2194806, upload-time = "2026-07-29T19:25:24.017Z" }, + { url = "https://files.pythonhosted.org/packages/cb/5b/03f4cddce426d005e208b39ea7b2d456e667cdee0f1891360f0fc8430f20/libcst-1.9.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:ea490fa8540503db5f321f0268becab46eb50f710e8cec8041e241fd66f6874f", size = 2246762, upload-time = "2026-07-29T19:25:25.365Z" }, + { url = "https://files.pythonhosted.org/packages/d8/31/9d5fe1e43dc3dbcc74f70f3e0e73fcdd8d84effc1059d8b45974f0d0d2eb/libcst-1.9.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:50ab94bb2524b419056d4003032b8c67102ac800f8d85b3c4260a01746d94dbb", size = 2259633, upload-time = "2026-07-29T19:25:26.928Z" }, + { url = "https://files.pythonhosted.org/packages/a8/2b/6752b28d88c3a19b3bc0b9e1838443b6760d5c862b4f4b37955402659e24/libcst-1.9.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:a2faaf92500d0226358125630f5aab4758e8aad3f2d70a10892ec3c700781a54", size = 2368043, upload-time = "2026-07-29T19:25:28.344Z" }, + { url = "https://files.pythonhosted.org/packages/e2/94/775825b2637f8ab05694b6a4b3802ae6783b4e799f9b58d2400c7e2d4369/libcst-1.9.0-cp314-cp314t-win_amd64.whl", hash = "sha256:0c7b548512db25af9c2997a95fa731bd6b6928ecbad6c0915d7482d8bb42d34f", size = 2176432, upload-time = "2026-07-29T19:25:29.921Z" }, + { url = "https://files.pythonhosted.org/packages/fb/3d/88ad67427c6fd9db929e087912b0a540e5140e5cb77e7ca4170edaac8531/libcst-1.9.0-cp314-cp314t-win_arm64.whl", hash = "sha256:497d5329345f1f5df84e41b0bbd00204b64a2fd30dfe3cfaeaebca633a31e877", size = 2052931, upload-time = "2026-07-29T19:25:31.397Z" }, + { url = "https://files.pythonhosted.org/packages/39/e4/ad790b043a38b10cea13166c8dd716654ad8b227c20f12eebf6326191cd9/libcst-1.9.0-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:a5068bf6114f6f4d79af7a6c80a28d1deb50441150ea1db13b5458ab34bec159", size = 2043987, upload-time = "2026-08-11T05:54:02.122Z" }, + { url = "https://files.pythonhosted.org/packages/15/6b/d3cd8275cc54ffc9817ade91442b3e6e9edd1a4518bd4e6dda13e3152dbe/libcst-1.9.0-cp315-cp315-manylinux_2_28_aarch64.whl", hash = "sha256:7acfd18adcdd32dcf41ce676cf121002afcbe9ad2c72dc0c18d78465beedc228", size = 2204003, upload-time = "2026-08-11T05:54:04.031Z" }, + { url = "https://files.pythonhosted.org/packages/da/62/87eddccb11d5d221d50ac4fff4b6e9b8d48c547028d01f7d0fbc5c8a1d75/libcst-1.9.0-cp315-cp315-manylinux_2_28_x86_64.whl", hash = "sha256:86361e2b426bd1b403e52375703c8b764e226e571adcb30442510710681edbf0", size = 2256053, upload-time = "2026-08-11T05:54:06.228Z" }, + { url = "https://files.pythonhosted.org/packages/18/53/a44126aeb9fca8b4e342e0cc803ea00ca8f6cd5712619a7897b3eba13797/libcst-1.9.0-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:8c14abe844bec8b021111b3985474e49f5af799c020e8e40dedcac87c97a40c5", size = 2270576, upload-time = "2026-08-11T05:54:08.01Z" }, + { url = "https://files.pythonhosted.org/packages/bb/51/3af3dd44b117d66486e0ed61e43a16b7fc8cc5c372e1dd4ca0e70b888d46/libcst-1.9.0-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:8930971d2299b7006bb5d10c10038f44efb6da89d5bca2823595e31a9cdb96af", size = 2378573, upload-time = "2026-08-11T05:54:10.093Z" }, + { url = "https://files.pythonhosted.org/packages/fc/e8/e545087d298dbf6494e84a8092f0f5dbd62eedacabd4ecd6b364367b4804/libcst-1.9.0-cp315-cp315-win_amd64.whl", hash = "sha256:5891ce9cff815077614f509e3a23888c5ddb8081d6188e8ba1172c0ccc369022", size = 2177937, upload-time = "2026-08-11T05:54:12.168Z" }, + { url = "https://files.pythonhosted.org/packages/84/17/93a00a8e03494a84db102dd0e3d35b8876b1084ca2c194b331a168d86eb1/libcst-1.9.0-cp315-cp315-win_arm64.whl", hash = "sha256:1260b5d070a3324447a53a00387225a55472a62077ce01922d30a2a78cc4b138", size = 2058633, upload-time = "2026-08-11T05:54:13.817Z" }, + { url = "https://files.pythonhosted.org/packages/d7/00/c3751b12eb81822ea0b6131d374a98bc6c024c4b9ec5f6ea8f53dc23e967/libcst-1.9.0-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:02ee2dbdb5c218116f16a021350fc267a14be205b50d81c33f5d73857ab6fbf7", size = 2036178, upload-time = "2026-08-11T05:54:15.545Z" }, + { url = "https://files.pythonhosted.org/packages/2f/12/48ff486eb5adc593f4818569e8896b74b672414ab4722a74b17e4c4cf31e/libcst-1.9.0-cp315-cp315t-manylinux_2_28_aarch64.whl", hash = "sha256:3e5b684191a0462b2d40a73261ea7f4da6a49e7a030b7e0fceca46bebb51dfe5", size = 2195496, upload-time = "2026-08-11T05:54:17.625Z" }, + { url = "https://files.pythonhosted.org/packages/51/18/b13a4669864d41a4801fed7b86ede1049dc9a1e1dde250a7ee1f9c3b1285/libcst-1.9.0-cp315-cp315t-manylinux_2_28_x86_64.whl", hash = "sha256:2b150c4f298fe54eb0fe73abf71f57db74d070796140a8c206c9615b6861f01e", size = 2245769, upload-time = "2026-08-11T05:54:19.343Z" }, + { url = "https://files.pythonhosted.org/packages/c8/c7/cf2d46744825e84afbd7228fdeeea8d23cf6c4b2074253c27efb7705c653/libcst-1.9.0-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:50b913e37187f00a6fb88962009364e1cf1b1284aa1762304f34b52b972f2218", size = 2260576, upload-time = "2026-08-11T05:54:21.444Z" }, + { url = "https://files.pythonhosted.org/packages/d1/4d/5433d3d62250b2e4325938db101766bcab2a006819684713bcccb6259a88/libcst-1.9.0-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:9f92c75283030fd58fb7d6e560168a00381e0e3368ec8a026a3ec8a828a05f93", size = 2368709, upload-time = "2026-08-11T05:54:23.276Z" }, + { url = "https://files.pythonhosted.org/packages/87/39/9f0e1690727623f99489895db0e42048a3e1e8cea0627517c593e437fd83/libcst-1.9.0-cp315-cp315t-win_amd64.whl", hash = "sha256:4adf97bb1aff8039b0bd4991e2f838be6e4fad552821b26b71a5d1a653c2582f", size = 2177866, upload-time = "2026-08-11T05:54:25.144Z" }, + { url = "https://files.pythonhosted.org/packages/cd/79/9dc7811883e67ea771057e8937bc536eb3e77f3635fc5a93cd85b6fe1ea8/libcst-1.9.0-cp315-cp315t-win_arm64.whl", hash = "sha256:8f0dd08a5773d7051e105250b2a2c73d8065ea9b2d68e7e1bdc5244660e75f93", size = 2052908, upload-time = "2026-08-11T05:54:26.835Z" }, +] + [[package]] name = "librt" version = "0.16.0" @@ -1464,6 +1528,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, ] +[[package]] +name = "pyyaml-ft" +version = "8.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/eb/5a0d575de784f9a1f94e2b1288c6886f13f34185e13117ed530f32b6f8a8/pyyaml_ft-8.0.0.tar.gz", hash = "sha256:0c947dce03954c7b5d38869ed4878b2e6ff1d44b08a0d84dc83fdad205ae39ab", size = 141057, upload-time = "2025-06-10T15:32:15.613Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/ba/a067369fe61a2e57fb38732562927d5bae088c73cb9bb5438736a9555b29/pyyaml_ft-8.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8c1306282bc958bfda31237f900eb52c9bedf9b93a11f82e1aab004c9a5657a6", size = 187027, upload-time = "2025-06-10T15:31:48.722Z" }, + { url = "https://files.pythonhosted.org/packages/ad/c5/a3d2020ce5ccfc6aede0d45bcb870298652ac0cf199f67714d250e0cdf39/pyyaml_ft-8.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:30c5f1751625786c19de751e3130fc345ebcba6a86f6bddd6e1285342f4bbb69", size = 176146, upload-time = "2025-06-10T15:31:50.584Z" }, + { url = "https://files.pythonhosted.org/packages/e3/bb/23a9739291086ca0d3189eac7cd92b4d00e9fdc77d722ab610c35f9a82ba/pyyaml_ft-8.0.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3fa992481155ddda2e303fcc74c79c05eddcdbc907b888d3d9ce3ff3e2adcfb0", size = 746792, upload-time = "2025-06-10T15:31:52.304Z" }, + { url = "https://files.pythonhosted.org/packages/5f/c2/e8825f4ff725b7e560d62a3609e31d735318068e1079539ebfde397ea03e/pyyaml_ft-8.0.0-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cec6c92b4207004b62dfad1f0be321c9f04725e0f271c16247d8b39c3bf3ea42", size = 786772, upload-time = "2025-06-10T15:31:54.712Z" }, + { url = "https://files.pythonhosted.org/packages/35/be/58a4dcae8854f2fdca9b28d9495298fd5571a50d8430b1c3033ec95d2d0e/pyyaml_ft-8.0.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:06237267dbcab70d4c0e9436d8f719f04a51123f0ca2694c00dd4b68c338e40b", size = 778723, upload-time = "2025-06-10T15:31:56.093Z" }, + { url = "https://files.pythonhosted.org/packages/86/ed/fed0da92b5d5d7340a082e3802d84c6dc9d5fa142954404c41a544c1cb92/pyyaml_ft-8.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:8a7f332bc565817644cdb38ffe4739e44c3e18c55793f75dddb87630f03fc254", size = 758478, upload-time = "2025-06-10T15:31:58.314Z" }, + { url = "https://files.pythonhosted.org/packages/f0/69/ac02afe286275980ecb2dcdc0156617389b7e0c0a3fcdedf155c67be2b80/pyyaml_ft-8.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7d10175a746be65f6feb86224df5d6bc5c049ebf52b89a88cf1cd78af5a367a8", size = 799159, upload-time = "2025-06-10T15:31:59.675Z" }, + { url = "https://files.pythonhosted.org/packages/4e/ac/c492a9da2e39abdff4c3094ec54acac9747743f36428281fb186a03fab76/pyyaml_ft-8.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:58e1015098cf8d8aec82f360789c16283b88ca670fe4275ef6c48c5e30b22a96", size = 158779, upload-time = "2025-06-10T15:32:01.029Z" }, + { url = "https://files.pythonhosted.org/packages/5d/9b/41998df3298960d7c67653669f37710fa2d568a5fc933ea24a6df60acaf6/pyyaml_ft-8.0.0-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:e64fa5f3e2ceb790d50602b2fd4ec37abbd760a8c778e46354df647e7c5a4ebb", size = 191331, upload-time = "2025-06-10T15:32:02.602Z" }, + { url = "https://files.pythonhosted.org/packages/0f/16/2710c252ee04cbd74d9562ebba709e5a284faeb8ada88fcda548c9191b47/pyyaml_ft-8.0.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:8d445bf6ea16bb93c37b42fdacfb2f94c8e92a79ba9e12768c96ecde867046d1", size = 182879, upload-time = "2025-06-10T15:32:04.466Z" }, + { url = "https://files.pythonhosted.org/packages/9a/40/ae8163519d937fa7bfa457b6f78439cc6831a7c2b170e4f612f7eda71815/pyyaml_ft-8.0.0-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8c56bb46b4fda34cbb92a9446a841da3982cdde6ea13de3fbd80db7eeeab8b49", size = 811277, upload-time = "2025-06-10T15:32:06.214Z" }, + { url = "https://files.pythonhosted.org/packages/f9/66/28d82dbff7f87b96f0eeac79b7d972a96b4980c1e445eb6a857ba91eda00/pyyaml_ft-8.0.0-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dab0abb46eb1780da486f022dce034b952c8ae40753627b27a626d803926483b", size = 831650, upload-time = "2025-06-10T15:32:08.076Z" }, + { url = "https://files.pythonhosted.org/packages/e8/df/161c4566facac7d75a9e182295c223060373d4116dead9cc53a265de60b9/pyyaml_ft-8.0.0-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bd48d639cab5ca50ad957b6dd632c7dd3ac02a1abe0e8196a3c24a52f5db3f7a", size = 815755, upload-time = "2025-06-10T15:32:09.435Z" }, + { url = "https://files.pythonhosted.org/packages/05/10/f42c48fa5153204f42eaa945e8d1fd7c10d6296841dcb2447bf7da1be5c4/pyyaml_ft-8.0.0-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:052561b89d5b2a8e1289f326d060e794c21fa068aa11255fe71d65baf18a632e", size = 810403, upload-time = "2025-06-10T15:32:11.051Z" }, + { url = "https://files.pythonhosted.org/packages/d5/d2/e369064aa51009eb9245399fd8ad2c562bd0bcd392a00be44b2a824ded7c/pyyaml_ft-8.0.0-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:3bb4b927929b0cb162fb1605392a321e3333e48ce616cdcfa04a839271373255", size = 835581, upload-time = "2025-06-10T15:32:12.897Z" }, + { url = "https://files.pythonhosted.org/packages/c0/28/26534bed77109632a956977f60d8519049f545abc39215d086e33a61f1f2/pyyaml_ft-8.0.0-cp313-cp313t-win_amd64.whl", hash = "sha256:de04cfe9439565e32f178106c51dd6ca61afaa2907d143835d501d84703d3793", size = 171579, upload-time = "2025-06-10T15:32:14.34Z" }, +] + [[package]] name = "pyzmq" version = "27.2.0" From c9ec72239a36901ee6661aa4edbd51ac57dc4940 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 05:01:54 +0200 Subject: [PATCH 10/17] style: rewrite src imports to module-only form --- src/array_vectorize/api.py | 28 ++- src/array_vectorize/codegen/docstring.py | 6 +- src/array_vectorize/codegen/emitter.py | 73 +++---- src/array_vectorize/compat.py | 4 +- src/array_vectorize/emit.py | 4 +- src/array_vectorize/errors.py | 4 +- src/array_vectorize/fallback.py | 20 +- src/array_vectorize/frontend/extract.py | 62 +++--- src/array_vectorize/frontend/info.py | 18 +- src/array_vectorize/frontend/lambda_id.py | 12 +- src/array_vectorize/frontend/validate.py | 17 +- src/array_vectorize/fuzz.py | 14 +- src/array_vectorize/ir/nodes.py | 28 +-- src/array_vectorize/ir/walk.py | 64 +++--- src/array_vectorize/lower/env.py | 85 ++++---- src/array_vectorize/lower/expressions.py | 187 +++++++++--------- src/array_vectorize/lower/helpers.py | 14 +- src/array_vectorize/lower/kinds.py | 34 ++-- src/array_vectorize/lower/loops.py | 48 ++--- src/array_vectorize/lower/sanitize.py | 50 ++--- src/array_vectorize/lower/statements.py | 26 +-- src/array_vectorize/lower/types.py | 12 +- src/array_vectorize/optimize/constfold.py | 91 ++++----- src/array_vectorize/optimize/cse.py | 49 +++-- src/array_vectorize/optimize/dce.py | 26 +-- .../optimize/protect_domains.py | 107 +++++----- src/array_vectorize/pipeline.py | 33 ++-- src/array_vectorize/runtime/arith.py | 6 +- src/array_vectorize/runtime/minmax.py | 23 +-- src/array_vectorize/runtime/registry.py | 11 +- src/array_vectorize/verify.py | 19 +- tests/features/test_namespace_pin.py | 116 ++++++----- 32 files changed, 620 insertions(+), 671 deletions(-) diff --git a/src/array_vectorize/api.py b/src/array_vectorize/api.py index ea5f4a9..3159bce 100644 --- a/src/array_vectorize/api.py +++ b/src/array_vectorize/api.py @@ -20,12 +20,10 @@ from collections.abc import Callable from typing import Any -from array_vectorize.compat import _check_namespace -from array_vectorize.errors import VectorizationError -from array_vectorize.fallback import make_fallback -from array_vectorize.frontend.extract import resolve_original -from array_vectorize.pipeline import compile_function -from array_vectorize.verify import verify_match +from array_vectorize import compat, errors, pipeline +from array_vectorize import fallback as fallback_mod +from array_vectorize import verify as verify_mod +from array_vectorize.frontend import extract __all__ = ["get_source", "vectorize"] @@ -52,7 +50,7 @@ def _vectorize_strict( namespace: Any = None, ) -> Callable[..., Any]: """Run the pipeline with helper vectorization wired to the api cache.""" - return compile_function( + return pipeline.compile_function( func, protect=protect, verify_args=verify_args, @@ -73,7 +71,7 @@ def _vectorize_helper( if key in _HELPER_CACHE: return _HELPER_CACHE[key] if callee in _ACTIVE_HELPERS: - raise VectorizationError( + raise errors.VectorizationError( f"cannot vectorize {callee.__name__!r}: recursive helper calls are not supported" ) _ACTIVE_HELPERS.add(callee) @@ -96,7 +94,7 @@ def _attach_namespace_api(vec: Callable[..., Any], *, protect: bool) -> None: """ def with_namespace(xp: Any) -> Callable[..., Any]: - _check_namespace(xp) + compat._check_namespace(xp) # chains to the true scalar original, so pins compose original: Callable[..., Any] = vec._vectorized_original # type: ignore[attr-defined] key = (original, protect, xp) @@ -166,8 +164,8 @@ def vectorize( but are rebuilt (and re-warn) on every call. """ if namespace is not None: - _check_namespace(namespace) # usage error: fail fast, before compiling - original = resolve_original(func) + compat._check_namespace(namespace) # usage error: fail fast, before compiling + original = extract.resolve_original(func) canonical = not protect_domains and namespace is None cached: Callable[..., Any] | None = None if canonical: @@ -184,14 +182,14 @@ def vectorize( # it (the user may pass different example args on each call). # The comparison target is the resolved scalar original — the # callable the cached function was compiled from. - verify_match(cached, original, verify) + verify_mod.verify_match(cached, original, verify) _attach_namespace_api(cached, protect=False) return cached try: vec = _vectorize_strict( func, protect=protect_domains, verify_args=verify, namespace=namespace ) - except VectorizationError as exc: + except errors.VectorizationError as exc: if fallback or not strict: reason = str(exc).splitlines()[0] warnings.warn( @@ -200,7 +198,7 @@ def vectorize( UserWarning, stacklevel=2, ) - wrapper = make_fallback(func, reason, namespace=namespace) + wrapper = fallback_mod.make_fallback(func, reason, namespace=namespace) if canonical: _set_backref(wrapper, original) return wrapper @@ -217,5 +215,5 @@ def get_source(vectorized: Callable[..., Any]) -> str: """Return the generated source of a vectorized function.""" source = getattr(vectorized, "source", None) if not isinstance(source, str): - raise VectorizationError(f"{vectorized!r} is not a vectorized function") + raise errors.VectorizationError(f"{vectorized!r} is not a vectorized function") return source diff --git a/src/array_vectorize/codegen/docstring.py b/src/array_vectorize/codegen/docstring.py index a8896ed..d087397 100644 --- a/src/array_vectorize/codegen/docstring.py +++ b/src/array_vectorize/codegen/docstring.py @@ -14,7 +14,7 @@ from __future__ import annotations -from inspect import cleandoc +import inspect __all__ = [ "DOCSTRING_PREFIX", @@ -144,7 +144,7 @@ def prefixed_summary(docstring: str | None) -> str: on already-cleaned input). Used by the fallback wrapper, which has only the original's ``__doc__`` available. """ - doc = cleandoc(docstring or "") + doc = inspect.cleandoc(docstring or "") if not doc: return DOCSTRING_PREFIX + NO_DOCSTRING_SUMMARY return _prefixed_head(doc) @@ -226,7 +226,7 @@ def build_docstring(docstring: str | None, source: str) -> str: synthesized line), and :func:`extract_scalar_source` recovers ``source`` from it exactly. """ - doc = cleandoc(docstring or "") + doc = inspect.cleandoc(docstring or "") head = _prefixed_head(doc) if doc else DOCSTRING_PREFIX + NO_DOCSTRING_SUMMARY if not doc: return head + "\n\n" + _new_google_notes(source) diff --git a/src/array_vectorize/codegen/emitter.py b/src/array_vectorize/codegen/emitter.py index 953f48c..45c8961 100644 --- a/src/array_vectorize/codegen/emitter.py +++ b/src/array_vectorize/codegen/emitter.py @@ -18,30 +18,13 @@ from __future__ import annotations import ast +import inspect import math -from inspect import cleandoc from typing import assert_never -from array_vectorize.codegen.docstring import build_docstring, extract_scalar_source -from array_vectorize.ir import ( - Binding, - BinOp, - Call, - Compare, - DType, - FuncCall, - Literal, - Logical, - Loop, - Node, - Program, - Ref, - Stmt, - UnaryOp, - Where, - generated_name, -) -from array_vectorize.lower.types import LoweredFunction +from array_vectorize import ir +from array_vectorize.codegen import docstring as docstring_mod +from array_vectorize.lower import types __all__ = ["generate_source"] @@ -84,7 +67,7 @@ def _xp_call(ns: str, attr: str, args: list[ast.expr]) -> ast.Call: return ast.Call(func=_xp_attr(ns, attr), args=args, keywords=[]) -def _gen_literal(lit: Literal, ns: str) -> ast.expr: +def _gen_literal(lit: ir.Literal, ns: str) -> ast.expr: value = lit.value if isinstance(value, float): if math.isnan(value): @@ -96,15 +79,15 @@ def _gen_literal(lit: Literal, ns: str) -> ast.expr: return ast.Constant(value=value) -def _gen_expr(node: Node, ns: str) -> ast.expr: +def _gen_expr(node: ir.Node, ns: str) -> ast.expr: match node: - case Literal(): + case ir.Literal(): return _gen_literal(node, ns) - case Ref(): + case ir.Ref(): return _load(node.name) - case DType(): + case ir.DType(): return _xp_attr(ns, node.name) - case BinOp(): + case ir.BinOp(): if node.op in _BINOP_AST: return ast.BinOp( op=_BINOP_AST[node.op](), @@ -114,7 +97,7 @@ def _gen_expr(node: Node, ns: str) -> ast.expr: return _xp_call( ns, _BINOP_XP[node.op], [_gen_expr(node.left, ns), _gen_expr(node.right, ns)] ) - case UnaryOp(): + case ir.UnaryOp(): if node.op == "neg": return ast.UnaryOp(op=ast.USub(), operand=_gen_expr(node.operand, ns)) if node.op == "pos": @@ -122,27 +105,27 @@ def _gen_expr(node: Node, ns: str) -> ast.expr: if node.op == "invert": return ast.UnaryOp(op=ast.Invert(), operand=_gen_expr(node.operand, ns)) return _xp_call(ns, "logical_not", [_gen_expr(node.operand, ns)]) - case Compare(): + case ir.Compare(): return ast.Compare( left=_gen_expr(node.left, ns), ops=[_CMPOP_AST[node.op]()], comparators=[_gen_expr(node.right, ns)], ) - case Logical(): + case ir.Logical(): fn = "logical_and" if node.op == "and" else "logical_or" folded = _gen_expr(node.parts[0], ns) for part in node.parts[1:]: folded = _xp_call(ns, fn, [folded, _gen_expr(part, ns)]) return folded - case Where(): + case ir.Where(): return _xp_call( ns, "where", [_gen_expr(node.cond, ns), _gen_expr(node.then, ns), _gen_expr(node.other, ns)], ) - case Call(): + case ir.Call(): return _xp_call(ns, node.fn, [_gen_expr(a, ns) for a in node.args]) - case FuncCall(): + case ir.FuncCall(): return ast.Call( func=_load(node.fn), args=[_gen_expr(a, ns) for a in node.args], keywords=[] ) @@ -164,7 +147,9 @@ def _namespace_line(ns: str, param_names: list[str]) -> ast.Assign: return ast.Assign(targets=[ast.Name(id=ns, ctx=ast.Store())], value=call) -def _build_signature(lowered: LoweredFunction, namespace_param: str | None = None) -> ast.arguments: +def _build_signature( + lowered: types.LoweredFunction, namespace_param: str | None = None +) -> ast.arguments: posonly: list[ast.arg] = [] positional: list[ast.arg] = [] defaults: list[ast.expr] = [] @@ -207,22 +192,22 @@ def _build_signature(lowered: LoweredFunction, namespace_param: str | None = Non ) -def _gen_range_call(loop: Loop, ns: str) -> ast.Call: - step_one = isinstance(loop.step, Literal) and loop.step.value == 1 +def _gen_range_call(loop: ir.Loop, ns: str) -> ast.Call: + step_one = isinstance(loop.step, ir.Literal) and loop.step.value == 1 args = [_gen_expr(loop.start, ns), _gen_expr(loop.stop, ns)] if not step_one: args.append(_gen_expr(loop.step, ns)) return ast.Call(func=_load("range"), args=args, keywords=[]) -def _gen_stmt(stmt: Stmt, ns: str) -> ast.stmt: +def _gen_stmt(stmt: ir.Stmt, ns: str) -> ast.stmt: match stmt: - case Binding(): + case ir.Binding(): return ast.Assign( targets=[ast.Name(id=stmt.name, ctx=ast.Store())], value=_gen_expr(stmt.expr, ns), ) - case Loop(): + case ir.Loop(): body = [_gen_stmt(inner, ns) for inner in stmt.body] if not body: # (15) a retained loop with a fully dead body still needs a statement @@ -252,10 +237,10 @@ def _indent_to_body(text: str) -> str: def _docstring_value(docstring: str | None, source: str) -> str: """Body-level docstring constant: combined format, or legacy on mismatch.""" - value = _indent_to_body(build_docstring(docstring, source)) + value = _indent_to_body(docstring_mod.build_docstring(docstring, source)) try: # simulate exactly what ast.get_docstring / inspect.getdoc return - if extract_scalar_source(cleandoc(value)) == source: + if docstring_mod.extract_scalar_source(inspect.cleandoc(value)) == source: return value except ValueError: pass @@ -264,7 +249,9 @@ def _docstring_value(docstring: str | None, source: str) -> str: return _indent_to_body(source) -def generate_source(lowered: LoweredFunction, program: Program, *, pinned: bool = False) -> str: +def generate_source( + lowered: types.LoweredFunction, program: ir.Program, *, pinned: bool = False +) -> str: """Generate the vectorized function source. With ``pinned=True`` (``vectorize(namespace=...)``) the namespace is @@ -272,7 +259,7 @@ def generate_source(lowered: LoweredFunction, program: Program, *, pinned: bool keyword-only ``namespace_param``, whose real default is injected at runtime, and the ``array_namespace`` import is omitted. """ - func_name = generated_name(lowered.name) + func_name = ir.generated_name(lowered.name) ns = lowered.namespace_var all_params = lowered.param_names + [name for name, _ in lowered.hidden_params] diff --git a/src/array_vectorize/compat.py b/src/array_vectorize/compat.py index dea8b73..9ea549a 100644 --- a/src/array_vectorize/compat.py +++ b/src/array_vectorize/compat.py @@ -16,7 +16,7 @@ from typing import Any -from array_api_compat import array_namespace +import array_api_compat __all__ = ["_check_namespace", "_is_array", "_py_scalar"] @@ -50,7 +50,7 @@ def _is_array(value: Any) -> bool: checks that miss backends. """ try: - array_namespace(value) + array_api_compat.array_namespace(value) except TypeError: return False return True diff --git a/src/array_vectorize/emit.py b/src/array_vectorize/emit.py index 66a529e..a0263b0 100644 --- a/src/array_vectorize/emit.py +++ b/src/array_vectorize/emit.py @@ -17,7 +17,7 @@ from collections.abc import Callable from typing import Any, cast -from array_vectorize.ir.ssa import generated_name +from array_vectorize.ir import ssa __all__ = ["compile_vectorized"] @@ -44,7 +44,7 @@ def compile_vectorized( exec(code, namespace) # retrieve the entry point BEFORE injecting helpers, so a helper whose # generated name collides can never shadow the function itself - func: Any = namespace[generated_name(name)] + func: Any = namespace[ssa.generated_name(name)] # helper functions (vectorized user helpers AND the runtime dtype # promotion helpers) are injected here; bare-name calls resolve through # the function's globals at call time diff --git a/src/array_vectorize/errors.py b/src/array_vectorize/errors.py index 659c35b..4403ffb 100644 --- a/src/array_vectorize/errors.py +++ b/src/array_vectorize/errors.py @@ -2,11 +2,11 @@ from __future__ import annotations -from dataclasses import dataclass +import dataclasses from typing import NoReturn -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Diagnostic: """One unsupported-construct finding with its source location.""" diff --git a/src/array_vectorize/fallback.py b/src/array_vectorize/fallback.py index 4163df7..79087e9 100644 --- a/src/array_vectorize/fallback.py +++ b/src/array_vectorize/fallback.py @@ -13,10 +13,10 @@ from collections.abc import Callable from typing import Any -from array_api_compat import array_namespace +import array_api_compat -from array_vectorize.codegen.docstring import prefixed_summary -from array_vectorize.compat import _check_namespace, _is_array, _py_scalar +from array_vectorize import compat +from array_vectorize.codegen import docstring __all__ = ["make_fallback"] @@ -32,15 +32,15 @@ def make_fallback( @functools.wraps(func) def wrapper(*args: Any, **kwargs: Any) -> Any: - positions = [i for i, a in enumerate(args) if _is_array(a)] - kw_positions = [k for k, v in kwargs.items() if _is_array(v)] + positions = [i for i, a in enumerate(args) if compat._is_array(a)] + kw_positions = [k for k, v in kwargs.items() if compat._is_array(v)] if not positions and not kw_positions: raise TypeError( f"vectorized fallback for {func.__name__!r} requires at least one " "Array API array argument" ) arrays = [args[i] for i in positions] + [kwargs[k] for k in kw_positions] - xp = namespace if namespace is not None else array_namespace(*arrays) + xp = namespace if namespace is not None else array_api_compat.array_namespace(*arrays) broadcast = xp.broadcast_arrays(*arrays) shape = broadcast[0].shape flat = [xp.reshape(a, (-1,)) for a in broadcast] @@ -50,14 +50,14 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: call_kwargs = dict(kwargs) values = iter(elems) for pos in positions: - call_args[pos] = _py_scalar(next(values)) + call_args[pos] = compat._py_scalar(next(values)) for key in kw_positions: - call_kwargs[key] = _py_scalar(next(values)) + call_kwargs[key] = compat._py_scalar(next(values)) out.append(func(*call_args, **call_kwargs)) return xp.reshape(xp.asarray(out), shape) def with_namespace(xp: Any) -> Callable[..., Any]: - _check_namespace(xp) + compat._check_namespace(xp) # cheap to build: no caching; the new wrapper gets its own # with_namespace, so pins chain return make_fallback(func, reason, namespace=xp) @@ -69,7 +69,7 @@ def with_namespace(xp: Any) -> Callable[..., Any]: # it is often unavailable here, which is a common fallback reason). # A pathological non-str __doc__ is left exactly as wraps left it. if isinstance(func.__doc__, str): - wrapper.__doc__ = prefixed_summary(func.__doc__) + wrapper.__doc__ = docstring.prefixed_summary(func.__doc__) wrapper.source = ( # type: ignore[attr-defined] f"# fallback element-loop wrapper around {func.__name__!r}\n" f"# (source-to-source vectorization failed: {reason})" diff --git a/src/array_vectorize/frontend/extract.py b/src/array_vectorize/frontend/extract.py index 85effb7..f9d6e45 100644 --- a/src/array_vectorize/frontend/extract.py +++ b/src/array_vectorize/frontend/extract.py @@ -11,15 +11,21 @@ from collections.abc import Callable from typing import Any -from array_vectorize.compat import _is_array -from array_vectorize.errors import _reject -from array_vectorize.frontend.info import FunctionInfo, Param, ParamKind, _AstFunction -from array_vectorize.frontend.lambda_id import _find_target -from array_vectorize.frontend.tables import MATH_FUNCS, MATH_SPECIAL +from array_vectorize import compat, errors +from array_vectorize.frontend import info as info_mod +from array_vectorize.frontend import lambda_id, tables + +# re-export for the module's public surface: __all__ names these, so the +# from-import is what keeps them importable from frontend.extract +from array_vectorize.frontend.info import ( # cleanporter: ignore[CP001] + FunctionInfo, + Param, + ParamKind, +) __all__ = ["FunctionInfo", "Param", "ParamKind", "extract_function", "resolve_original"] -_MATH_FUNC_NAMES = (*MATH_FUNCS, *MATH_SPECIAL) +_MATH_FUNC_NAMES = (*tables.MATH_FUNCS, *tables.MATH_SPECIAL) def resolve_original(func: Callable[..., Any]) -> Callable[..., Any]: @@ -40,36 +46,40 @@ def resolve_original(func: Callable[..., Any]) -> Callable[..., Any]: def _check_default(param: str, node: ast.expr) -> int | float | bool: if isinstance(node, ast.Constant) and type(node.value) in (int, float, bool): return node.value # type: ignore[return-value] - _reject(param, f"default values must be int/float/bool literals (got {ast.unparse(node)!r})") + errors._reject( + param, f"default values must be int/float/bool literals (got {ast.unparse(node)!r})" + ) -def _params_of(fn: _AstFunction, name: str) -> list[Param]: +def _params_of(fn: info_mod._AstFunction, name: str) -> list[info_mod.Param]: a = fn.args if a.vararg is not None or a.kwarg is not None: - _reject(name, "*args/**kwargs are not supported") - params: list[Param] = [] + errors._reject(name, "*args/**kwargs are not supported") + params: list[info_mod.Param] = [] positional = [*a.posonlyargs, *a.args] defaults: dict[int, ast.expr] = {} offset = len(positional) - len(a.defaults) for i, default in enumerate(a.defaults): defaults[offset + i] = default for i, arg in enumerate(positional): - kind: ParamKind = "posonly" if i < len(a.posonlyargs) else "arg" + kind: info_mod.ParamKind = "posonly" if i < len(a.posonlyargs) else "arg" if i in defaults: - params.append(Param(arg.arg, kind, _check_default(arg.arg, defaults[i]), True)) + params.append(info_mod.Param(arg.arg, kind, _check_default(arg.arg, defaults[i]), True)) else: - params.append(Param(arg.arg, kind, None, False)) + params.append(info_mod.Param(arg.arg, kind, None, False)) for arg, kw_default in zip(a.kwonlyargs, a.kw_defaults, strict=True): if kw_default is not None: - params.append(Param(arg.arg, "kwonly", _check_default(arg.arg, kw_default), True)) + params.append( + info_mod.Param(arg.arg, "kwonly", _check_default(arg.arg, kw_default), True) + ) else: - params.append(Param(arg.arg, "kwonly", None, False)) + params.append(info_mod.Param(arg.arg, "kwonly", None, False)) if not params: - _reject(name, "zero-argument functions cannot be vectorized") + errors._reject(name, "zero-argument functions cannot be vectorized") return params -def _classify_closures(target: types.FunctionType, info: FunctionInfo) -> None: +def _classify_closures(target: types.FunctionType, info: info_mod.FunctionInfo) -> None: # getclosurevars works off co_names, which includes attribute names # (``math.exp`` contributes ``exp``); filter to actual Name references. loaded = {n.id for n in ast.walk(info.tree) if isinstance(n, ast.Name)} @@ -79,7 +89,7 @@ def _classify_closures(target: types.FunctionType, info: FunctionInfo) -> None: unbound = cv.unbound & loaded if unbound: names = ", ".join(sorted(unbound)) - _reject(info.name, f"name(s) not resolvable in the function's scope: {names}") + errors._reject(info.name, f"name(s) not resolvable in the function's scope: {names}") for var, value in resolved.items(): if value is math: info.math_modules.add(var) @@ -88,12 +98,12 @@ def _classify_closures(target: types.FunctionType, info: FunctionInfo) -> None: info.math_funcs[var] = attr elif isinstance(value, bool | int | float): info.closure_scalars[var] = value - elif _is_array(value): + elif compat._is_array(value): info.closure_arrays[var] = value elif isinstance(value, types.FunctionType): info.user_funcs[var] = value else: - _reject( + errors._reject( info.name, f"closure/global {var!r} has unsupported type {type(value).__name__!r}; " "supported: int/float/bool scalars, Array API arrays, math functions, " @@ -101,16 +111,16 @@ def _classify_closures(target: types.FunctionType, info: FunctionInfo) -> None: ) -def extract_function(func: Callable[..., Any]) -> FunctionInfo: +def extract_function(func: Callable[..., Any]) -> info_mod.FunctionInfo: """Unwrap, read source, parse, and capture closures.""" target = resolve_original(func) if not isinstance(target, types.FunctionType): if inspect.isbuiltin(target): - _reject( + errors._reject( getattr(target, "__name__", repr(target)), "C-implemented builtins have no inspectable Python source", ) - _reject( + errors._reject( getattr(target, "__name__", repr(target)), f"expected a plain Python function, got {type(target).__name__}", ) @@ -118,7 +128,7 @@ def extract_function(func: Callable[..., Any]) -> FunctionInfo: try: raw = inspect.getsource(target) except (OSError, TypeError): - _reject( + errors._reject( name, "source is not available (defined in a REPL or via exec()?); " "define the function in a .py file so inspect.getsource works", @@ -131,7 +141,7 @@ def extract_function(func: Callable[..., Any]) -> FunctionInfo: module_source = None if filename and filename != "": module_source = "".join(linecache.getlines(filename)) - tree = _find_target( + tree = lambda_id._find_target( ast.parse(textwrap.dedent(raw)), name, target, @@ -142,7 +152,7 @@ def extract_function(func: Callable[..., Any]) -> FunctionInfo: params = _params_of(tree, name) docstring = ast.get_docstring(tree) if isinstance(tree, ast.FunctionDef) else None - info = FunctionInfo( + info = info_mod.FunctionInfo( name=name, filename=filename or "", source=textwrap.dedent(raw).strip(), diff --git a/src/array_vectorize/frontend/info.py b/src/array_vectorize/frontend/info.py index f0a646a..eae7325 100644 --- a/src/array_vectorize/frontend/info.py +++ b/src/array_vectorize/frontend/info.py @@ -3,8 +3,8 @@ from __future__ import annotations import ast +import dataclasses from collections.abc import Callable -from dataclasses import dataclass, field from typing import Any, Literal __all__ = ["FunctionInfo", "Param", "ParamKind"] @@ -15,7 +15,7 @@ type ParamKind = Literal["posonly", "arg", "kwonly"] -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Param: """One function parameter; ``kind`` mirrors Python argument kinds.""" @@ -25,7 +25,7 @@ class Param: has_default: bool = False -@dataclass(slots=True) +@dataclasses.dataclass(slots=True) class FunctionInfo: """Everything the later passes need about the scalar function.""" @@ -35,12 +35,12 @@ class FunctionInfo: tree: _AstFunction params: list[Param] docstring: str | None - closure_scalars: dict[str, int | float | bool] = field(default_factory=dict) - closure_arrays: dict[str, Any] = field(default_factory=dict) - math_funcs: dict[str, str] = field(default_factory=dict) # name -> math attr - math_modules: set[str] = field(default_factory=set) # names bound to `math` - user_funcs: dict[str, Callable[..., Any]] = field(default_factory=dict) - user_names: set[str] = field(default_factory=set) + closure_scalars: dict[str, int | float | bool] = dataclasses.field(default_factory=dict) + closure_arrays: dict[str, Any] = dataclasses.field(default_factory=dict) + math_funcs: dict[str, str] = dataclasses.field(default_factory=dict) # name -> math attr + math_modules: set[str] = dataclasses.field(default_factory=set) # names bound to `math` + user_funcs: dict[str, Callable[..., Any]] = dataclasses.field(default_factory=dict) + user_names: set[str] = dataclasses.field(default_factory=set) @property def param_names(self) -> list[str]: diff --git a/src/array_vectorize/frontend/lambda_id.py b/src/array_vectorize/frontend/lambda_id.py index 8a5e318..f7b9270 100644 --- a/src/array_vectorize/frontend/lambda_id.py +++ b/src/array_vectorize/frontend/lambda_id.py @@ -14,8 +14,8 @@ from collections.abc import Iterator from typing import Any -from array_vectorize.errors import _reject -from array_vectorize.frontend.info import _AstFunction +from array_vectorize import errors +from array_vectorize.frontend import info def _lambdas_in_source_order(tree: ast.AST) -> Iterator[ast.Lambda]: @@ -87,10 +87,10 @@ def _find_target( module_source: str | None = None, first_lineno: int | None = None, filename: str | None = None, -) -> _AstFunction: +) -> info._AstFunction: for stmt in tree.body: if isinstance(stmt, ast.AsyncFunctionDef): - _reject(name, "async functions are not supported") + errors._reject(name, "async functions are not supported") if isinstance(stmt, ast.FunctionDef | ast.Lambda): return stmt if ( @@ -114,7 +114,7 @@ def _find_target( # compare bytecode, constants, referenced names, and default values. candidates = list(_lambdas_in_source_order(tree)) if not candidates: - _reject(name, "could not locate the function definition in its source") + errors._reject(name, "could not locate the function definition in its source") if len(candidates) == 1: return candidates[0] code = getattr(target, "__code__", None) @@ -168,7 +168,7 @@ def _find_target( if len(matches) == 1: return matches[0] if not matches: - _reject( + errors._reject( name, "multiple lambdas share this source line and the intended one " "could not be identified; assign the lambda to a variable on its " diff --git a/src/array_vectorize/frontend/validate.py b/src/array_vectorize/frontend/validate.py index 7216e9a..255d029 100644 --- a/src/array_vectorize/frontend/validate.py +++ b/src/array_vectorize/frontend/validate.py @@ -9,9 +9,8 @@ import ast -from array_vectorize.errors import Diagnostic, VectorizationError -from array_vectorize.frontend.info import FunctionInfo -from array_vectorize.frontend.tables import MATH_CONSTS, MATH_FUNCS, MATH_SPECIAL +from array_vectorize import errors +from array_vectorize.frontend import info, tables __all__ = ["validate"] @@ -32,7 +31,7 @@ _ALLOWED_UNARY = (ast.USub, ast.UAdd, ast.Invert, ast.Not) _ALLOWED_CMPOPS = (ast.Eq, ast.NotEq, ast.Lt, ast.LtE, ast.Gt, ast.GtE) -_MATH_ATTRS = set(MATH_FUNCS) | set(MATH_SPECIAL) | set(MATH_CONSTS) +_MATH_ATTRS = set(tables.MATH_FUNCS) | set(tables.MATH_SPECIAL) | set(tables.MATH_CONSTS) _STMT_MESSAGES: dict[type[ast.stmt], str] = { ast.While: "while loops are not supported (data-dependent control flow)", @@ -87,17 +86,17 @@ class _Validator: - def __init__(self, info: FunctionInfo) -> None: + def __init__(self, info: info.FunctionInfo) -> None: self.info = info self.lines = info.source.splitlines() - self.diags: list[Diagnostic] = [] + self.diags: list[errors.Diagnostic] = [] self.saw_return = False def error(self, node: ast.AST, message: str) -> None: lineno = getattr(node, "lineno", 1) col = getattr(node, "col_offset", 0) line = self.lines[lineno - 1] if 0 < lineno <= len(self.lines) else "" - self.diags.append(Diagnostic(message, lineno, col, line)) + self.diags.append(errors.Diagnostic(message, lineno, col, line)) # ------------------------------------------------------------- statements @@ -277,7 +276,7 @@ def validate_call(self, node: ast.Call) -> None: self.validate_expr(arg) -def validate(info: FunctionInfo) -> None: +def validate(info: info.FunctionInfo) -> None: """Validate the supported subset; raise with all violations if any.""" v = _Validator(info) tree = info.tree @@ -293,7 +292,7 @@ def validate(info: FunctionInfo) -> None: v.validate_expr(tree.body) if v.diags: n = len(v.diags) - raise VectorizationError( + raise errors.VectorizationError( f"cannot vectorize {info.name!r}: {n} unsupported construct" f"{'' if n == 1 else 's'} found", v.diags, diff --git a/src/array_vectorize/fuzz.py b/src/array_vectorize/fuzz.py index 07fe268..25ffcce 100644 --- a/src/array_vectorize/fuzz.py +++ b/src/array_vectorize/fuzz.py @@ -11,18 +11,22 @@ import argparse import importlib.util +import pathlib import random import shutil import sys import tempfile import time -from pathlib import Path from typing import Any import numpy as np -from array_vectorize import vectorize -from array_vectorize.errors import VectorizationError +# load-bearing re-export: tests/unit/test_fuzz.py monkeypatches +# array_vectorize.fuzz.vectorize, so the name must stay bound here +from array_vectorize import ( + errors, + vectorize, # cleanporter: ignore[CP003] load-bearing +) _UNARY_MATH = [ f"math.{name}" @@ -139,7 +143,7 @@ def make_program(rng: random.Random) -> str: def _run_one(rng: random.Random) -> str: src = make_program(rng) # vectorize() requires inspectable source: define the case in a real file - case_dir = Path(tempfile.mkdtemp(prefix="array_vectorize_fuzz_")) + case_dir = pathlib.Path(tempfile.mkdtemp(prefix="array_vectorize_fuzz_")) path = case_dir / "fuzz_case_mod.py" path.write_text("import math\n\n\n" + src) spec = importlib.util.spec_from_file_location("fuzz_case_mod", path) @@ -158,7 +162,7 @@ def _run_one(rng: random.Random) -> str: def _check_case(rng: random.Random, src: str, fn: Any) -> str: try: vec = vectorize(fn) - except VectorizationError: + except errors.VectorizationError: return "rejected" except Exception as exc: raise AssertionError(f"internal error vectorizing:\n{src}\n{exc}") from exc diff --git a/src/array_vectorize/ir/nodes.py b/src/array_vectorize/ir/nodes.py index 33294b8..426e2ca 100644 --- a/src/array_vectorize/ir/nodes.py +++ b/src/array_vectorize/ir/nodes.py @@ -7,9 +7,9 @@ from __future__ import annotations +import dataclasses import math import typing -from dataclasses import dataclass from typing import Any __all__ = [ @@ -45,7 +45,7 @@ def _literal_key(value: int | float | bool) -> tuple[Any, ...]: type Kind = typing.Literal["int", "float", "bool"] -@dataclass(frozen=True, eq=False, slots=True) +@dataclasses.dataclass(frozen=True, eq=False, slots=True) class Literal: """A scalar constant. ``inf``/``nan`` allowed (codegen emits xp.inf / xp.nan). @@ -68,14 +68,14 @@ def __hash__(self) -> int: return hash((self.kind, self.value)) -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Ref: """Reference to a parameter or SSA binding by its emitted name.""" name: str -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class BinOp: """Binary operation, e.g. ``BinOp('add', x, y)`` -> x + y.""" @@ -84,7 +84,7 @@ class BinOp: right: Node -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class UnaryOp: """Unary operation, e.g. ``UnaryOp('neg', x)`` -> -x.""" @@ -92,7 +92,7 @@ class UnaryOp: operand: Node -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Compare: """Comparison, e.g. ``Compare('lt', x, y)`` -> x < y.""" @@ -101,7 +101,7 @@ class Compare: right: Node -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Logical: """``and``/``or`` over bool-typed operands only (numeric and/or lowers to Where).""" @@ -109,7 +109,7 @@ class Logical: parts: tuple[Node, ...] -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Where: """Branch merge: ``then`` where ``cond`` holds, ``other`` elsewhere.""" @@ -118,7 +118,7 @@ class Where: other: Node -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Call: """Call to an Array API function, e.g. ``Call('sqrt', (x,))`` -> xp.sqrt(x).""" @@ -126,14 +126,14 @@ class Call: args: tuple[Node, ...] -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class DType: """An xp dtype reference; ``DType('int64')`` codegens to ``xp.int64``.""" name: str -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class FuncCall: """Call to a vectorized helper function by its emitted name (design D7).""" @@ -144,7 +144,7 @@ class FuncCall: type Node = Literal | Ref | BinOp | UnaryOp | Compare | Logical | Where | Call | DType | FuncCall -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Binding: """Let-statement: ``name = expr``.""" @@ -152,7 +152,7 @@ class Binding: expr: Node -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Loop: """Constant-trip ``for`` loop emitted as a real loop (design D6). @@ -171,7 +171,7 @@ class Loop: type Stmt = Binding | Loop -@dataclass(frozen=True, slots=True) +@dataclasses.dataclass(frozen=True, slots=True) class Program: """A lowered function: parameters, top-level bindings, result.""" diff --git a/src/array_vectorize/ir/walk.py b/src/array_vectorize/ir/walk.py index e5c7447..366c6ff 100644 --- a/src/array_vectorize/ir/walk.py +++ b/src/array_vectorize/ir/walk.py @@ -9,64 +9,52 @@ from collections.abc import Callable from typing import assert_never -from array_vectorize.ir.nodes import ( - BinOp, - Call, - Compare, - DType, - FuncCall, - Literal, - Logical, - Node, - Ref, - UnaryOp, - Where, -) +from array_vectorize.ir import nodes __all__ = ["children", "rewrite"] -def rewrite(node: Node, fn: Callable[[Node], Node]) -> Node: +def rewrite(node: nodes.Node, fn: Callable[[nodes.Node], nodes.Node]) -> nodes.Node: """Rebuild ``node`` with ``fn`` applied to each rewritten child.""" match node: - case Literal() | Ref() | DType(): + case nodes.Literal() | nodes.Ref() | nodes.DType(): return node - case BinOp(): - return BinOp(node.op, fn(node.left), fn(node.right)) - case UnaryOp(): - return UnaryOp(node.op, fn(node.operand)) - case Compare(): - return Compare(node.op, fn(node.left), fn(node.right)) - case Logical(): - return Logical(node.op, tuple(fn(p) for p in node.parts)) - case Where(): - return Where(fn(node.cond), fn(node.then), fn(node.other)) - case Call(): - return Call(node.fn, tuple(fn(a) for a in node.args)) - case FuncCall(): - return FuncCall(node.fn, tuple(fn(a) for a in node.args)) + case nodes.BinOp(): + return nodes.BinOp(node.op, fn(node.left), fn(node.right)) + case nodes.UnaryOp(): + return nodes.UnaryOp(node.op, fn(node.operand)) + case nodes.Compare(): + return nodes.Compare(node.op, fn(node.left), fn(node.right)) + case nodes.Logical(): + return nodes.Logical(node.op, tuple(fn(p) for p in node.parts)) + case nodes.Where(): + return nodes.Where(fn(node.cond), fn(node.then), fn(node.other)) + case nodes.Call(): + return nodes.Call(node.fn, tuple(fn(a) for a in node.args)) + case nodes.FuncCall(): + return nodes.FuncCall(node.fn, tuple(fn(a) for a in node.args)) case _: assert_never(node) -def children(node: Node) -> tuple[Node, ...]: +def children(node: nodes.Node) -> tuple[nodes.Node, ...]: """The direct operand sub-nodes of ``node`` (leaves have none).""" match node: - case Literal() | Ref() | DType(): + case nodes.Literal() | nodes.Ref() | nodes.DType(): return () - case BinOp(): + case nodes.BinOp(): return (node.left, node.right) - case UnaryOp(): + case nodes.UnaryOp(): return (node.operand,) - case Compare(): + case nodes.Compare(): return (node.left, node.right) - case Logical(): + case nodes.Logical(): return node.parts - case Where(): + case nodes.Where(): return (node.cond, node.then, node.other) - case Call(): + case nodes.Call(): return node.args - case FuncCall(): + case nodes.FuncCall(): return node.args case _: assert_never(node) diff --git a/src/array_vectorize/lower/env.py b/src/array_vectorize/lower/env.py index 9606ed2..1c05e3c 100644 --- a/src/array_vectorize/lower/env.py +++ b/src/array_vectorize/lower/env.py @@ -14,23 +14,10 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any -from array_vectorize.errors import VectorizationError -from array_vectorize.frontend.info import FunctionInfo -from array_vectorize.ir import ( - Binding, - Compare, - Kind, - Literal, - Node, - Ref, - SSAEnv, - Stmt, - Where, - is_bool, -) -from array_vectorize.lower.kinds import Kinds -from array_vectorize.lower.types import HelperVectorizer -from array_vectorize.runtime.registry import RUNTIME_HELPERS +from array_vectorize import errors, ir +from array_vectorize.frontend import info +from array_vectorize.lower import kinds, types +from array_vectorize.runtime import registry __all__ = ["_LowererBase"] @@ -44,29 +31,31 @@ class _LowererBase: # detection, and helper dispatch from above this base). Type-only # declarations — never executed; the real methods on the assembled # class are the ones that run. - def lower_stmts(self, stmts: list[ast.stmt], *, loop_body: bool = False) -> Node | None: ... - def _numeric_kind(self, node: Node) -> Kind | None: ... - def _maybe_bool_result(self, node: Node) -> bool: ... - def _call_helper(self, node: ast.Call, name: str, args: list[Node]) -> Node: ... + def lower_stmts( + self, stmts: list[ast.stmt], *, loop_body: bool = False + ) -> ir.Node | None: ... + def _numeric_kind(self, node: ir.Node) -> ir.Kind | None: ... + def _maybe_bool_result(self, node: ir.Node) -> bool: ... + def _call_helper(self, node: ast.Call, name: str, args: list[ir.Node]) -> ir.Node: ... def __init__( self, - info: FunctionInfo, - helper_vectorizer: HelperVectorizer | None = None, + info: info.FunctionInfo, + helper_vectorizer: types.HelperVectorizer | None = None, ) -> None: self.info = info self.helper_vectorizer = helper_vectorizer self.helpers: list[tuple[str, Any]] = [] self._helper_names: dict[Callable[..., Any], str] = {} - self.ssa = SSAEnv(info.user_names) - self.bindings: list[Stmt] = [] + self.ssa = ir.SSAEnv(info.user_names) + self.bindings: list[ir.Stmt] = [] #: definite locals: var -> current emitted name self.definite: dict[str, str] = {} #: vars that may be unbound on the current path (assigned in one branch only) self.maybe: set[str] = set() #: pending early returns: (condition, value, push_depth) in divergence #: order; folds nest earlier entries outermost - self.deferred: list[tuple[Node, Node, int]] = [] + self.deferred: list[tuple[ir.Node, ir.Node, int]] = [] #: nesting depth inside branches (branch-local bindings always suffix) self.branch_depth = 0 #: nesting depth inside loop bodies (returns are rejected there) @@ -77,14 +66,14 @@ def __init__( self.active_carried: list[tuple[dict[str, str], int]] = [] #: kind/literal/scalar/bool bookkeeping for emitted names, plus the #: per-loop carried-assign kind frames (see kinds.Kinds) - self.kinds = Kinds() + self.kinds = kinds.Kinds() #: the generated namespace variable ('xp' unless taken); set by #: lower_function before lowering starts self.ns_var: str = "xp" #: allocated names of the runtime promotion/selection helpers (set #: by lower_function; collision-free against user names) - self.minmax_name = RUNTIME_HELPERS["vec_minmax"][0] - self.arith_name = RUNTIME_HELPERS["vec_arith"][0] + self.minmax_name = registry.RUNTIME_HELPERS["vec_minmax"][0] + self.arith_name = registry.RUNTIME_HELPERS["vec_arith"][0] self.assigned_names: set[str] = { t.id for stmt in ast.walk(info.tree) @@ -101,29 +90,31 @@ def __init__( # ------------------------------------------------------------- utilities - def error(self, node: ast.AST, message: str) -> VectorizationError: + def error(self, node: ast.AST, message: str) -> errors.VectorizationError: lineno = getattr(node, "lineno", 0) - return VectorizationError(f"cannot vectorize {self.info.name!r}: {message} (line {lineno})") + return errors.VectorizationError( + f"cannot vectorize {self.info.name!r}: {message} (line {lineno})" + ) - def _literal(self, value: Any) -> Literal: + def _literal(self, value: Any) -> ir.Literal: if isinstance(value, bool): - return Literal(value, "bool") + return ir.Literal(value, "bool") if isinstance(value, int): - return Literal(int(value), "int") - return Literal(float(value), "float") + return ir.Literal(int(value), "int") + return ir.Literal(float(value), "float") - def _coerce_bool(self, node: Node) -> Node: + def _coerce_bool(self, node: ir.Node) -> ir.Node: """Coerce a possibly-numeric condition to bool via != 0 (design D2/D4).""" - if is_bool(node): + if ir.is_bool(node): return node - return Compare("ne", node, Literal(0, "int")) + return ir.Compare("ne", node, ir.Literal(0, "int")) # ---------------------------------------------------------- name access - def load_name(self, node: ast.Name) -> Node: + def load_name(self, node: ast.Name) -> ir.Node: var = node.id if var in self.definite: - return Ref(self.definite[var]) + return ir.Ref(self.definite[var]) if var in self.maybe: raise self.error(node, f"variable {var!r} may be unbound (assigned on some paths only)") if var in self.assigned_names: @@ -132,7 +123,7 @@ def load_name(self, node: ast.Name) -> Node: if var in info.closure_scalars: return self._literal(info.closure_scalars[var]) if var in info.closure_arrays: - return Ref(self.definite[var]) # hidden param, bound at intake + return ir.Ref(self.definite[var]) # hidden param, bound at intake if var in info.math_funcs or var in info.user_funcs or var in info.math_modules: raise self.error(node, f"cannot reference {var!r} without calling it") raise self.error(node, f"name {var!r} is not supported in vectorized code") @@ -151,7 +142,7 @@ def _mark_then_only_unbound( def _merge_envs( self, - cond: Node, + cond: ir.Node, pre: dict[str, str], then_def: dict[str, str], else_def: dict[str, str], @@ -164,8 +155,8 @@ def _merge_envs( continue if nt is not None and ne is not None: name = self.ssa.bind(var) - merged_expr: Node = Where(cond, Ref(nt), Ref(ne)) - self.bindings.append(Binding(name, merged_expr)) + merged_expr: ir.Node = ir.Where(cond, ir.Ref(nt), ir.Ref(ne)) + self.bindings.append(ir.Binding(name, merged_expr)) self.definite[var] = name kt, ke = self.kinds.kind(nt), self.kinds.kind(ne) self.kinds.set_kind(name, kt if kt == ke else None) @@ -183,7 +174,7 @@ def _merge_envs( self.maybe.discard(var) self.definite[var] = then_def[var] - def fold_returns(self, value: Node) -> Node: + def fold_returns(self, value: ir.Node) -> ir.Node: """Fold pending early returns around ``value``. Entries pushed at the current branch depth or deeper belong to this @@ -194,12 +185,12 @@ def fold_returns(self, value: Node) -> Node: divergence before enclosed entries. """ result = value - retained: list[tuple[Node, Node, int]] = [] + retained: list[tuple[ir.Node, ir.Node, int]] = [] for cond, early, depth in reversed(self.deferred): if depth < self.branch_depth: retained.append((cond, early, depth)) else: - result = Where(cond, early, result) + result = ir.Where(cond, early, result) retained.reverse() self.deferred = retained return result diff --git a/src/array_vectorize/lower/expressions.py b/src/array_vectorize/lower/expressions.py index 2b153d1..e6a31ef 100644 --- a/src/array_vectorize/lower/expressions.py +++ b/src/array_vectorize/lower/expressions.py @@ -12,32 +12,10 @@ import ast -from array_vectorize.frontend.tables import ( - BUILTIN_CASTS, - BUILTIN_FOLDS, - BUILTIN_UNARY, - MATH_CONSTS, - MATH_FUNCS, - MATH_SPECIAL, -) -from array_vectorize.ir import ( - Binding, - BinOp, - Call, - Compare, - DType, - FuncCall, - Kind, - Literal, - Logical, - Node, - Ref, - UnaryOp, - Where, - is_bool, -) -from array_vectorize.lower.sanitize import _Sanitizer -from array_vectorize.runtime.registry import ArithOp +from array_vectorize import ir +from array_vectorize.frontend import tables +from array_vectorize.lower import sanitize +from array_vectorize.runtime import registry __all__ = ["_ExpressionLowerer"] @@ -71,10 +49,10 @@ } -class _ExpressionLowerer(_Sanitizer): +class _ExpressionLowerer(sanitize._Sanitizer): # --------------------------------------------------------- expressions - def lower_expr(self, node: ast.expr) -> Node: + def lower_expr(self, node: ast.expr) -> ir.Node: match node: case ast.Constant(): return self._literal(node.value) @@ -92,19 +70,23 @@ def lower_expr(self, node: ast.expr) -> Node: # -True is -1 in Python; numpy cannot negate bool arrays. # The whole negation goes through the runtime helper (it # computes the negation — no wrapping UnaryOp on top) - if isinstance(operand, Literal): - return Literal(-int(operand.value), "int") - return FuncCall( + if isinstance(operand, ir.Literal): + return ir.Literal(-int(operand.value), "int") + return ir.FuncCall( self.arith_name, - (Ref(self.ns_var), Literal(int(ArithOp.NEG), "int"), operand), + ( + ir.Ref(self.ns_var), + ir.Literal(int(registry.ArithOp.NEG), "int"), + operand, + ), ) - return UnaryOp(_UNARY[type(node.op)], operand) + return ir.UnaryOp(_UNARY[type(node.op)], operand) case ast.Compare(): return self._lower_compare(node) case ast.BoolOp(): return self._lower_boolop(node) case ast.IfExp(): - return Where( + return ir.Where( self._coerce_bool(self.lower_expr(node.test)), self.lower_expr(node.body), self.lower_expr(node.orelse), @@ -118,41 +100,43 @@ def lower_expr(self, node: ast.expr) -> Node: # the validator rejects these first, lowering never sees them raise self.error(node, f"{type(node).__name__} expressions are not supported") - def _lower_compare(self, node: ast.Compare) -> Node: - parts: list[Node] = [] + def _lower_compare(self, node: ast.Compare) -> ir.Node: + parts: list[ir.Node] = [] operands = [node.left, *node.comparators] for left, op, right in zip(operands, node.ops, operands[1:], strict=False): - parts.append(Compare(_CMPOPS[type(op)], self.lower_expr(left), self.lower_expr(right))) + parts.append( + ir.Compare(_CMPOPS[type(op)], self.lower_expr(left), self.lower_expr(right)) + ) if len(parts) == 1: return parts[0] - return Logical("and", tuple(parts)) + return ir.Logical("and", tuple(parts)) - def _lower_boolop(self, node: ast.BoolOp) -> Node: + def _lower_boolop(self, node: ast.BoolOp) -> ir.Node: op = "and" if isinstance(node.op, ast.And) else "or" values = [self.lower_expr(v) for v in node.values] - if all(is_bool(v) for v in values): - return Logical(op, tuple(values)) + if all(ir.is_bool(v) for v in values): + return ir.Logical(op, tuple(values)) # exact value-select lowering (design D2) partial = values[0] last = len(values) - 1 for i, value in enumerate(values[1:], 1): if op == "and": - merged: Node = Where(self._coerce_bool(partial), value, partial) + merged: ir.Node = ir.Where(self._coerce_bool(partial), value, partial) else: - merged = Where(self._coerce_bool(partial), partial, value) + merged = ir.Where(self._coerce_bool(partial), partial, value) if i < last: # bind intermediates to temps; only the final value stays inline temp = self.ssa.fresh_temp("v") - self.bindings.append(Binding(temp, merged)) - partial = Ref(temp) + self.bindings.append(ir.Binding(temp, merged)) + partial = ir.Ref(temp) else: partial = merged return partial - def _lower_math_const(self, node: ast.Attribute) -> Node: - return self._literal(MATH_CONSTS[node.attr]) + def _lower_math_const(self, node: ast.Attribute) -> ir.Node: + return self._literal(tables.MATH_CONSTS[node.attr]) - def _lower_call(self, node: ast.Call) -> Node: + def _lower_call(self, node: ast.Call) -> ir.Node: args = [self.lower_expr(a) for a in node.args] func = node.func if isinstance(func, ast.Attribute): @@ -173,18 +157,18 @@ def _lower_call(self, node: ast.Call) -> Node: node, "round(x, n) is not supported: the decimals kwarg is not in the Array API standard", ) - if name in BUILTIN_UNARY: + if name in tables.BUILTIN_UNARY: self._check_arity(node, name, args, 1, 1) sanitized = tuple(self._sanitize_xp_arg(a, force_float=False) for a in args) - return Call(BUILTIN_UNARY[name], sanitized) - if name in BUILTIN_FOLDS: + return ir.Call(tables.BUILTIN_UNARY[name], sanitized) + if name in tables.BUILTIN_FOLDS: if len(args) < 2: raise self.error(node, f"{name}() needs at least 2 arguments in vectorized code") if all(self._literal_value(a) is not None for a in args): # all-literal call: substitute the actual literals so # const-folding fires (Ref nodes would block it and reach # codegen as raw Python scalars) - clean: list[Node] = [ + clean: list[ir.Node] = [ lit for a in args if (lit := self._literal_value(a)) is not None ] else: @@ -194,25 +178,25 @@ def _lower_call(self, node: ast.Call) -> Node: # uint64; min(uint64, signed) fits int64), Python float # promotion. Literal args pass as raw values; array args # pass through the soft sanitizer - is_min = BUILTIN_FOLDS[name] == "minimum" - helper_args: list[Node] = [Ref(self.ns_var), Literal(is_min, "bool")] + is_min = tables.BUILTIN_FOLDS[name] == "minimum" + helper_args: list[ir.Node] = [ir.Ref(self.ns_var), ir.Literal(is_min, "bool")] for a in args: lit = self._literal_value(a) if lit is None: helper_args.append(self._sanitize_xp_arg(a, force_float=False)) else: helper_args.append(lit) - return FuncCall(self.minmax_name, tuple(helper_args)) - folded: Node = Call(BUILTIN_FOLDS[name], (clean[0], clean[1])) + return ir.FuncCall(self.minmax_name, tuple(helper_args)) + folded: ir.Node = ir.Call(tables.BUILTIN_FOLDS[name], (clean[0], clean[1])) for arg in clean[2:]: - folded = Call(BUILTIN_FOLDS[name], (folded, arg)) + folded = ir.Call(tables.BUILTIN_FOLDS[name], (folded, arg)) return folded - if name in BUILTIN_CASTS: + if name in tables.BUILTIN_CASTS: self._check_arity(node, name, args, 1, 1) - return self._cast_call(args[0], BUILTIN_CASTS[name]) + return self._cast_call(args[0], tables.BUILTIN_CASTS[name]) raise self.error(node, f"unsupported function {name!r}") - def _cast_call(self, arg: Node, dtype: str) -> Node: + def _cast_call(self, arg: ir.Node, dtype: str) -> ir.Node: """Lower an int()/float()/bool()/math.trunc-style cast. Literal args fold exactly, raw loop-variable ints wrap in asarray, @@ -222,59 +206,59 @@ def _cast_call(self, arg: Node, dtype: str) -> Node: if lit is not None: try: if dtype == "bool": - return Literal(bool(lit.value), "bool") + return ir.Literal(bool(lit.value), "bool") if dtype == "int64": - return Literal(int(lit.value), "int") - return Literal(float(lit.value), "float") + return ir.Literal(int(lit.value), "int") + return ir.Literal(float(lit.value), "float") except (ValueError, OverflowError): pass # inf/nan literals: leave the cast to runtime astype - return Call("astype", (self._sanitize_xp_arg(arg, force_float=False), DType(dtype))) + return ir.Call("astype", (self._sanitize_xp_arg(arg, force_float=False), ir.DType(dtype))) - def _call_math(self, func: ast.AST, attr: str, args: list[Node], node: ast.AST) -> Node: - if attr in MATH_SPECIAL: + def _call_math(self, func: ast.AST, attr: str, args: list[ir.Node], node: ast.AST) -> ir.Node: + if attr in tables.MATH_SPECIAL: self._check_arity(node, attr, args, 1, 1) - return self._cast_call(args[0], MATH_SPECIAL[attr]) - if attr in MATH_FUNCS: - xp_name, lo, hi = MATH_FUNCS[attr] + return self._cast_call(args[0], tables.MATH_SPECIAL[attr]) + if attr in tables.MATH_FUNCS: + xp_name, lo, hi = tables.MATH_FUNCS[attr] self._check_arity(node, attr, args, lo, hi) # Array API elementwise math requires floating-point array # inputs; sanitize literal/int/bool/possibly-scalar args sanitized = tuple(self._sanitize_xp_arg(a, force_float=True) for a in args) - return Call(xp_name, sanitized) + return ir.Call(xp_name, sanitized) raise self.error(func, f"math.{attr} cannot be called") @staticmethod - def _is_negative_int(node: Node) -> bool: + def _is_negative_int(node: ir.Node) -> bool: """True for literal negative ints (including UnaryOp(neg, lit)).""" - if isinstance(node, Literal) and node.kind == "int": + if isinstance(node, ir.Literal) and node.kind == "int": return node.value < 0 return ( - isinstance(node, UnaryOp) + isinstance(node, ir.UnaryOp) and node.op == "neg" - and isinstance(node.operand, Literal) + and isinstance(node.operand, ir.Literal) and node.operand.kind == "int" ) - def _numeric_kind(self, node: Node) -> Kind | None: + def _numeric_kind(self, node: ir.Node) -> ir.Kind | None: """Best-effort numeric kind of ``node``. Literals, casts, bools, arithmetic, and Refs whose binding kind was recorded are provable; anything else reads as ``None``. """ match node: - case Literal(): + case ir.Literal(): return node.kind - case Ref(): + case ir.Ref(): return self.kinds.kind(node.name) - case Logical(): + case ir.Logical(): return "bool" - case Where(): + case ir.Where(): kt, ke = self._numeric_kind(node.then), self._numeric_kind(node.other) return kt if kt == ke else None - case _ if is_bool(node): + case _ if ir.is_bool(node): # provably-boolean nodes (Compare, not, isnan/isinf/isfinite) return "bool" - case BinOp(): + case ir.BinOp(): # scalar promotion rules: int/bool + float -> float; arithmetic # over ints/bools stays int; true division is ALWAYS float in # Python (int / int -> float), whatever the operand kinds @@ -294,12 +278,12 @@ def _numeric_kind(self, node: Node) -> Kind | None: if lk in ("int", "bool") and rk in ("int", "bool"): return "int" return None - case UnaryOp() if node.op in ("neg", "pos"): + case ir.UnaryOp() if node.op in ("neg", "pos"): return self._numeric_kind(node.operand) - case Call() if node.fn == "asarray" and len(node.args) == 1: + case ir.Call() if node.fn == "asarray" and len(node.args) == 1: # asarray preserves the wrapped value's kind return self._numeric_kind(node.args[0]) - case FuncCall() if node.fn == self.minmax_name: + case ir.FuncCall() if node.fn == self.minmax_name: # min/max result kind: bool when every argument is bool, float # when any argument is float, int for int-ish mixes (the helper # normalizes bool arrays to int8, so arithmetic downstream must @@ -312,15 +296,15 @@ def _numeric_kind(self, node: Node) -> Kind | None: if all(k == "bool" for k in kinds): return "bool" return "int" - case FuncCall() if node.fn == self.arith_name: + case ir.FuncCall() if node.fn == self.arith_name: # arithmetic helper result kind: floats stay float, true # division is float, known int/bool operands give int (bools # are converted to a numeric dtype inside the helper) — so # downstream mitigations (e.g. the negative-exponent float # cast) keep firing for provably-integer results op_lit = node.args[1] if len(node.args) > 1 else None - op_id = op_lit.value if isinstance(op_lit, Literal) else None - if op_id == int(ArithOp.DIV): + op_id = op_lit.value if isinstance(op_lit, ir.Literal) else None + if op_id == int(registry.ArithOp.DIV): return "float" kinds = [self._numeric_kind(a) for a in node.args[2:]] if any(k is None for k in kinds): @@ -328,8 +312,8 @@ def _numeric_kind(self, node: Node) -> Kind | None: if "float" in kinds: return "float" return "int" - case Call() if ( - node.fn == "astype" and len(node.args) == 2 and isinstance(node.args[1], DType) + case ir.Call() if ( + node.fn == "astype" and len(node.args) == 2 and isinstance(node.args[1], ir.DType) ): name = node.args[1].name return "bool" if name == "bool" else "int" if name == "int64" else "float" @@ -337,7 +321,7 @@ def _numeric_kind(self, node: Node) -> Kind | None: # no provable kind (unknown ops, untracked casts, ...) return None - def _maybe_bool_result(self, node: Node) -> bool: + def _maybe_bool_result(self, node: ir.Node) -> bool: """True for min/max results whose static kind is unknown. Their runtime dtype may be boolean (parameter operands), so @@ -345,20 +329,20 @@ def _maybe_bool_result(self, node: Node) -> bool: _vec_arith_dtype cast is a no-op for numeric dtypes and int64 for booleans. """ - if isinstance(node, FuncCall) and node.fn == self.minmax_name: + if isinstance(node, ir.FuncCall) and node.fn == self.minmax_name: return self._numeric_kind(node) is None - if isinstance(node, Ref): + if isinstance(node, ir.Ref): return self.kinds.is_maybe_bool(node.name) - if isinstance(node, Where): + if isinstance(node, ir.Where): # branch merge of maybe-bool values is maybe-bool return self._maybe_bool_result(node.then) or self._maybe_bool_result(node.other) - if isinstance(node, BinOp) and node.op in ("and", "or", "xor"): + if isinstance(node, ir.BinOp) and node.op in ("and", "or", "xor"): # bitwise expressions preserve boolean-ness (arithmetic # results are already intified and numeric) return self._maybe_bool_result(node.left) or self._maybe_bool_result(node.right) return False - def _lower_binop(self, op: ast.operator, left: Node, right: Node) -> Node: + def _lower_binop(self, op: ast.operator, left: ir.Node, right: ir.Node) -> ir.Node: """Build a BinOp with the exactness fixes shared by `a op b` and `a op= b`.""" # Python bool arithmetic is integer (True + True == 2), but array # backends either saturate (bool + bool == True) or reject bools in @@ -377,17 +361,22 @@ def _lower_binop(self, op: ast.operator, left: Node, right: Node) -> Node: # uint64 integer paths are exact per lane (modular for # add/sub/mul; int64 magnitudes for negative divisors) ir_op = _BINOPS[type(op)] - return FuncCall( + return ir.FuncCall( self.arith_name, - (Ref(self.ns_var), Literal(int(ArithOp[ir_op.upper()]), "int"), left, right), + ( + ir.Ref(self.ns_var), + ir.Literal(int(registry.ArithOp[ir_op.upper()]), "int"), + left, + right, + ), ) if isinstance(op, ast.Pow) and self._is_negative_int(right): # Python promotes int ** negative-int to float; NumPy raises. # Cast provably-int bases so the common literal case matches. left = self._as_float_arg(left) - return BinOp(_BINOPS[type(op)], left, right) + return ir.BinOp(_BINOPS[type(op)], left, right) - def _check_arity(self, node: ast.AST, name: str, args: list[Node], lo: int, hi: int) -> None: + def _check_arity(self, node: ast.AST, name: str, args: list[ir.Node], lo: int, hi: int) -> None: if not lo <= len(args) <= hi: expected = str(lo) if lo == hi else f"{lo}-{hi}" raise self.error(node, f"{name}() takes {expected} argument(s), got {len(args)}") diff --git a/src/array_vectorize/lower/helpers.py b/src/array_vectorize/lower/helpers.py index e27aeec..b2de4d8 100644 --- a/src/array_vectorize/lower/helpers.py +++ b/src/array_vectorize/lower/helpers.py @@ -10,14 +10,14 @@ import ast -from array_vectorize.ir import FuncCall, Node, generated_name -from array_vectorize.lower.loops import _LoopLowerer +from array_vectorize import ir +from array_vectorize.lower import loops __all__ = ["_HelperLowerer"] -class _HelperLowerer(_LoopLowerer): - def _call_helper(self, node: ast.Call, name: str, args: list[Node]) -> Node: +class _HelperLowerer(loops._LoopLowerer): + def _call_helper(self, node: ast.Call, name: str, args: list[ir.Node]) -> ir.Node: """Vectorize and call another scalar function (design D7).""" if self.helper_vectorizer is None: raise self.error( @@ -27,12 +27,12 @@ def _call_helper(self, node: ast.Call, name: str, args: list[Node]) -> Node: ) callee = self.info.user_funcs[name] if callee in self._helper_names: - return FuncCall(self._helper_names[callee], tuple(args)) - base = generated_name(callee.__name__) + return ir.FuncCall(self._helper_names[callee], tuple(args)) + base = ir.generated_name(callee.__name__) while not self.ssa.is_free(base): base += "_" self.ssa.reserve(base) vec = self.helper_vectorizer(callee) self.helpers.append((base, vec)) self._helper_names[callee] = base - return FuncCall(base, tuple(args)) + return ir.FuncCall(base, tuple(args)) diff --git a/src/array_vectorize/lower/kinds.py b/src/array_vectorize/lower/kinds.py index ea78b32..ed2daac 100644 --- a/src/array_vectorize/lower/kinds.py +++ b/src/array_vectorize/lower/kinds.py @@ -42,25 +42,25 @@ from __future__ import annotations +import dataclasses from collections.abc import Mapping -from dataclasses import dataclass -from array_vectorize.ir import Kind, Literal +from array_vectorize import ir __all__ = ["FactsSnapshot", "Kinds", "VarInfo"] -@dataclass +@dataclasses.dataclass class VarInfo: """Per-name kind-inference facts (replaces the five parallel dicts).""" - kind: Kind | None = None - literal: Literal | None = None + kind: ir.Kind | None = None + literal: ir.Literal | None = None maybe_scalar: bool = False maybe_bool: bool = False -@dataclass(frozen=True) +@dataclasses.dataclass(frozen=True) class FactsSnapshot: """Presence-faithful copy of the kind + literal facts. @@ -70,8 +70,8 @@ class FactsSnapshot: so values are always real literals). """ - kinds: dict[str, Kind | None] - literals: dict[str, Literal] + kinds: dict[str, ir.Kind | None] + literals: dict[str, ir.Literal] class Kinds: @@ -87,7 +87,7 @@ def __init__(self) -> None: # "unknown" VALUE — presence is what setdefault consults) self._kind_established: set[str] = set() # per active loop: {carried name -> kinds assigned in the body} - self.carried_frames: list[dict[str, set[Kind | None]]] = [] + self.carried_frames: list[dict[str, set[ir.Kind | None]]] = [] def _entry(self, name: str) -> VarInfo: info = self._names.get(name) @@ -97,35 +97,35 @@ def _entry(self, name: str) -> VarInfo: # ------------------------------------------------------------ kind facts - def kind(self, name: str) -> Kind | None: + def kind(self, name: str) -> ir.Kind | None: """The name's recorded kind; absent names read as ``None``.""" if name not in self._kind_established: return None return self._names[name].kind - def set_kind(self, name: str, kind: Kind | None) -> None: + def set_kind(self, name: str, kind: ir.Kind | None) -> None: """Record ``name``'s kind, replacing any previous fact.""" self._kind_established.add(name) self._entry(name).kind = kind - def setdefault_kind(self, name: str, kind: Kind | None) -> None: + def setdefault_kind(self, name: str, kind: ir.Kind | None) -> None: """Insert-only: an established fact (even ``kind=None``) wins.""" if name not in self._kind_established: self.set_kind(name, kind) - def update_kinds(self, labels: Mapping[str, Kind | None]) -> None: + def update_kinds(self, labels: Mapping[str, ir.Kind | None]) -> None: """``set_kind`` for every name/kind entry at once.""" for name, kind in labels.items(): self.set_kind(name, kind) # --------------------------------------------------------- literal facts - def literal(self, name: str) -> Literal | None: + def literal(self, name: str) -> ir.Literal | None: """The name's recorded literal, or ``None``.""" info = self._names.get(name) return info.literal if info is not None else None - def set_literal(self, name: str, lit: Literal) -> None: + def set_literal(self, name: str, lit: ir.Literal) -> None: """Record ``name`` as holding the literal ``lit``.""" self._entry(name).literal = lit @@ -173,11 +173,11 @@ def push_carried_frame(self) -> None: """Enter a loop: body assignments record into a fresh frame.""" self.carried_frames.append({}) - def pop_carried_frame(self) -> dict[str, set[Kind | None]]: + def pop_carried_frame(self) -> dict[str, set[ir.Kind | None]]: """Leave a loop, returning its carried-assignment frame.""" return self.carried_frames.pop() - def record_carried_kind(self, name: str, kind: Kind | None) -> None: + def record_carried_kind(self, name: str, kind: ir.Kind | None) -> None: """Record a body-assignment kind for ``name`` in EVERY active frame.""" for frame in self.carried_frames: frame.setdefault(name, set()).add(kind) diff --git a/src/array_vectorize/lower/loops.py b/src/array_vectorize/lower/loops.py index ea0aa8b..9efabd2 100644 --- a/src/array_vectorize/lower/loops.py +++ b/src/array_vectorize/lower/loops.py @@ -12,39 +12,39 @@ import ast -from array_vectorize.ir import Binding, Kind, Literal, Loop, Node, Ref -from array_vectorize.lower.statements import _StatementLowerer -from array_vectorize.optimize.constfold import _const_fold_expr +from array_vectorize import ir +from array_vectorize.lower import statements +from array_vectorize.optimize import constfold __all__ = ["_LoopLowerer"] -class _LoopLowerer(_StatementLowerer): +class _LoopLowerer(statements._StatementLowerer): # ------------------------------------------------------------- loops - def _const_int(self, node: ast.expr, what: str) -> Node: - folded = _const_fold_expr(self.lower_expr(node)) + def _const_int(self, node: ast.expr, what: str) -> ir.Node: + folded = constfold._const_fold_expr(self.lower_expr(node)) if ( - isinstance(folded, Literal) + isinstance(folded, ir.Literal) and folded.kind == "int" and -(2**63) <= folded.value < 2**63 ): return folded raise self.error(node, f"{what} must be constant ints known at generation time") - def _loop_bounds(self, call: ast.Call) -> tuple[Node, Node, Node]: + def _loop_bounds(self, call: ast.Call) -> tuple[ir.Node, ir.Node, ir.Node]: args = list(call.args) - bounds: list[Node] = [self._const_int(a, "loop bound") for a in args] + bounds: list[ir.Node] = [self._const_int(a, "loop bound") for a in args] if len(bounds) == 1: - start_n: Node = Literal(0, "int") - stop_n: Node = bounds[0] - step_n: Node = Literal(1, "int") + start_n: ir.Node = ir.Literal(0, "int") + stop_n: ir.Node = bounds[0] + step_n: ir.Node = ir.Literal(1, "int") elif len(bounds) == 2: start_n, stop_n = bounds - step_n = Literal(1, "int") + step_n = ir.Literal(1, "int") else: start_n, stop_n, step_n = bounds - assert isinstance(step_n, Literal) + assert isinstance(step_n, ir.Literal) if step_n.value == 0: raise self.error(call, "range step cannot be zero") return start_n, stop_n, step_n @@ -106,18 +106,18 @@ def lower_for(self, stmt: ast.For) -> None: if loop_var in body_assigned: raise self.error(stmt, f"cannot assign the loop variable {loop_var!r} inside its loop") - phi_kinds: dict[str, Kind | None] = {} + phi_kinds: dict[str, ir.Kind | None] = {} # a pre-bound loop variable keeps its value on zero-trip loops: # emit a phi (the for-statement overwrites it on real iterations) if loop_var in pre_definite: - self.bindings.append(Binding(loop_name, Ref(pre_definite[loop_var]))) + self.bindings.append(ir.Binding(loop_name, ir.Ref(pre_definite[loop_var]))) self.kinds.setdefault_kind(loop_name, self.kinds.kind(pre_definite[loop_var])) # phis for loop-carried variables, emitted just before the loop for var, pre_name in carried.items(): loop_carried_name = self.ssa.bind(var) carried[var] = loop_carried_name - self.bindings.append(Binding(loop_carried_name, Ref(pre_name))) + self.bindings.append(ir.Binding(loop_carried_name, ir.Ref(pre_name))) self.definite[var] = loop_carried_name phi_kinds[loop_carried_name] = self.kinds.kind(pre_name) self.kinds.set_kind(loop_carried_name, self.kinds.kind(pre_name)) @@ -131,15 +131,15 @@ def lower_for(self, stmt: ast.For) -> None: phi_end = len(self.bindings) def union_labels( - frame: dict[str, set[Kind | None]], - ) -> tuple[dict[str, Kind | None], list[str]]: + frame: dict[str, set[ir.Kind | None]], + ) -> tuple[dict[str, ir.Kind | None], list[str]]: """Per carried name: the final kind and the mixed names. The final kind is the union of the phi kind and all body-assignment kinds; a name is mixed when its kinds MIX across iterations. """ - labels: dict[str, Kind | None] = {} + labels: dict[str, ir.Kind | None] = {} mixed: list[str] = [] for loop_carried_name in carried.values(): kinds = {phi_kinds.get(loop_carried_name)} | frame.get(loop_carried_name, set()) @@ -162,8 +162,8 @@ def union_labels( # cannot see every name that needs the intify conversion. Labels # only ever WIDEN (int/float -> bool), so at most len(carried) # widenings can happen; bound the passes at len(carried) + 2. - labels: dict[str, Kind | None] = {} - prev_labels: dict[str, Kind | None] | None = None + labels: dict[str, ir.Kind | None] = {} + prev_labels: dict[str, ir.Kind | None] | None = None max_passes = len(carried) + 2 for attempt in range(max_passes): pre_body = ( @@ -210,7 +210,7 @@ def union_labels( self.kinds.drop_literal(loop_carried_name) body_stmts = tuple(self.bindings[phi_end:]) del self.bindings[phi_end:] - self.bindings.append(Loop(loop_name, start, stop, step, body_stmts)) + self.bindings.append(ir.Loop(loop_name, start, stop, step, body_stmts)) # loop-local variables (assigned in body, not bound before) keep their # body names: zero-trip loops raise NameError exactly like Python. @@ -222,7 +222,7 @@ def _sync_loop_carried(self) -> None: for var, loop_name in carried.items(): current = self.definite.get(var) if current is not None and current != loop_name: - self.bindings.append(Binding(loop_name, Ref(current))) + self.bindings.append(ir.Binding(loop_name, ir.Ref(current))) self.definite[var] = loop_name self.kinds.set_kind(loop_name, self.kinds.kind(current)) self.kinds.drop_literal(loop_name) diff --git a/src/array_vectorize/lower/sanitize.py b/src/array_vectorize/lower/sanitize.py index c5bd3e1..c1d22ea 100644 --- a/src/array_vectorize/lower/sanitize.py +++ b/src/array_vectorize/lower/sanitize.py @@ -10,34 +10,34 @@ from __future__ import annotations -from array_vectorize.ir import Call, DType, Literal, Node, Ref, UnaryOp -from array_vectorize.ir.walk import children -from array_vectorize.lower.env import _LowererBase +from array_vectorize import ir +from array_vectorize.ir import walk as walk_mod +from array_vectorize.lower import env __all__ = ["_Sanitizer"] -class _Sanitizer(_LowererBase): - def _literal_value(self, arg: Node) -> Literal | None: +class _Sanitizer(env._LowererBase): + def _literal_value(self, arg: ir.Node) -> ir.Literal | None: """The literal behind an arg, or ``None``. Found directly, as a negated literal (``-1`` lowers to UnaryOp(neg, 1)), or via a literal binding. """ - if isinstance(arg, Literal): + if isinstance(arg, ir.Literal): return arg if ( - isinstance(arg, UnaryOp) + isinstance(arg, ir.UnaryOp) and arg.op == "neg" - and isinstance(arg.operand, Literal) + and isinstance(arg.operand, ir.Literal) and arg.operand.kind in ("int", "float") ): - return Literal(-arg.operand.value, arg.operand.kind) - if isinstance(arg, Ref): + return ir.Literal(-arg.operand.value, arg.operand.kind) + if isinstance(arg, ir.Ref): return self.kinds.literal(arg.name) return None - def _possibly_scalar(self, node: Node) -> bool: + def _possibly_scalar(self, node: ir.Node) -> bool: """True when the node's runtime value may be a raw Python scalar. It is a literal, a possibly-scalar name (parameter default, loop @@ -46,18 +46,18 @@ def _possibly_scalar(self, node: Node) -> bool: not (strict backends reject plain scalars). """ - def walk(n: Node) -> bool: - if isinstance(n, Literal): + def walk(n: ir.Node) -> bool: + if isinstance(n, ir.Literal): return True - if isinstance(n, Ref): + if isinstance(n, ir.Ref): return self.kinds.is_scalar(n.name) - if isinstance(n, DType): + if isinstance(n, ir.DType): return False - return any(walk(c) for c in children(n)) + return any(walk(c) for c in walk_mod.children(n)) return walk(node) - def _sanitize_xp_arg(self, arg: Node, *, force_float: bool) -> Node: + def _sanitize_xp_arg(self, arg: ir.Node, *, force_float: bool) -> ir.Node: """Make an argument safe for xp.* calls. Operators promote raw Python scalars fine, but xp.* function @@ -76,19 +76,19 @@ def _sanitize_xp_arg(self, arg: Node, *, force_float: bool) -> Node: lit = self._literal_value(arg) if lit is not None: if force_float: - return Literal(float(lit.value), "float") - return Call("asarray", (lit,)) + return ir.Literal(float(lit.value), "float") + return ir.Call("asarray", (lit,)) kind = self._numeric_kind(arg) if kind in ("int", "bool") or self._possibly_scalar(arg): if force_float: - return Call("astype", (Call("asarray", (arg,)), DType("float64"))) - return Call("asarray", (arg,)) + return ir.Call("astype", (ir.Call("asarray", (arg,)), ir.DType("float64"))) + return ir.Call("asarray", (arg,)) return arg - def _as_float_arg(self, arg: Node) -> Node: + def _as_float_arg(self, arg: ir.Node) -> ir.Node: kind = self._numeric_kind(arg) if kind in ("int", "bool"): - if isinstance(arg, Literal): - return Literal(float(arg.value), "float") - return Call("astype", (Call("asarray", (arg,)), DType("float64"))) + if isinstance(arg, ir.Literal): + return ir.Literal(float(arg.value), "float") + return ir.Call("astype", (ir.Call("asarray", (arg,)), ir.DType("float64"))) return arg diff --git a/src/array_vectorize/lower/statements.py b/src/array_vectorize/lower/statements.py index 046505e..aa86741 100644 --- a/src/array_vectorize/lower/statements.py +++ b/src/array_vectorize/lower/statements.py @@ -10,13 +10,13 @@ import ast -from array_vectorize.ir import Binding, Literal, Logical, Node, UnaryOp, Where, is_bool -from array_vectorize.lower.expressions import _ExpressionLowerer +from array_vectorize import ir +from array_vectorize.lower import expressions __all__ = ["_StatementLowerer"] -class _StatementLowerer(_ExpressionLowerer): +class _StatementLowerer(expressions._ExpressionLowerer): # --------------------------------------------------------- statements def lower_assign(self, stmt: ast.Assign | ast.AugAssign) -> None: @@ -42,11 +42,11 @@ def lower_assign(self, stmt: ast.Assign | ast.AugAssign) -> None: name = carried[var] if name is None: name = self.ssa.bind(var, force_suffix=self.branch_depth > 0) - self.bindings.append(Binding(name, value)) + self.bindings.append(ir.Binding(name, value)) self.definite[var] = name kind = self._numeric_kind(value) self.kinds.set_kind(name, kind) - if isinstance(value, Literal): + if isinstance(value, ir.Literal): self.kinds.set_literal(name, value) else: self.kinds.drop_literal(name) @@ -67,9 +67,9 @@ def lower_assign(self, stmt: ast.Assign | ast.AugAssign) -> None: self.kinds.record_carried_kind(carried[var], kind) self.maybe.discard(var) - def lower_if(self, stmt: ast.If) -> Node | None: + def lower_if(self, stmt: ast.If) -> ir.Node | None: cond = self.lower_expr(stmt.test) - if not is_bool(cond): + if not ir.is_bool(cond): raise self.error( stmt, "bare truthiness in an if condition is ambiguous per-lane; " @@ -102,7 +102,7 @@ def lower_if(self, stmt: ast.If) -> Node | None: # entries pending at this block's depth (they fire earlier in # program order, so they wrap the merged where from outside). self.deferred = list(pre_deferred) - merged: Node = Where(cond, then_result, else_result) + merged: ir.Node = ir.Where(cond, then_result, else_result) return self.fold_returns(merged) if then_result is not None: @@ -121,7 +121,7 @@ def lower_if(self, stmt: ast.If) -> Node | None: if else_result is not None: # else returns, then falls through: continue on the then path. self.definite, self.maybe = then_definite, then_maybe - not_cond = UnaryOp("not", cond) + not_cond = ir.UnaryOp("not", cond) self.deferred = [ *then_deferred[:pre_len], (not_cond, else_result, self.branch_depth), @@ -133,11 +133,11 @@ def lower_if(self, stmt: ast.If) -> Node | None: # both fall through: merge environments with where-selects self._merge_envs(cond, pre_definite, then_definite, else_definite, then_maybe, else_maybe) # qualify branch-pending returns with their branch conditions - not_cond = UnaryOp("not", cond) - qualified: list[tuple[Node, Node, int]] = [] + not_cond = ir.UnaryOp("not", cond) + qualified: list[tuple[ir.Node, ir.Node, int]] = [] for c, v, p in then_deferred[pre_len:]: - qualified.append((Logical("and", (cond, c)), v, p)) + qualified.append((ir.Logical("and", (cond, c)), v, p)) for c, v, p in else_deferred[pre_len:]: - qualified.append((Logical("and", (not_cond, c)), v, p)) + qualified.append((ir.Logical("and", (not_cond, c)), v, p)) self.deferred = [*pre_deferred, *qualified] return None diff --git a/src/array_vectorize/lower/types.py b/src/array_vectorize/lower/types.py index d552160..c874027 100644 --- a/src/array_vectorize/lower/types.py +++ b/src/array_vectorize/lower/types.py @@ -8,23 +8,23 @@ from __future__ import annotations +import dataclasses from collections.abc import Callable -from dataclasses import dataclass from typing import Any -from array_vectorize.frontend.info import Param -from array_vectorize.ir import Program +from array_vectorize import ir +from array_vectorize.frontend import info __all__ = ["HelperVectorizer", "LoweredFunction"] -@dataclass(slots=True) +@dataclasses.dataclass(slots=True) class LoweredFunction: """Lowering output: IR plus everything codegen/runtime need.""" - program: Program + program: ir.Program name: str - params: list[Param] + params: list[info.Param] param_names: list[str] # emitted names, aligned with params hidden_params: list[tuple[str, Any]] # (emitted name, array default) source: str # original scalar source (embedded verbatim in the docstring) diff --git a/src/array_vectorize/optimize/constfold.py b/src/array_vectorize/optimize/constfold.py index 0ce5cbe..fda3d87 100644 --- a/src/array_vectorize/optimize/constfold.py +++ b/src/array_vectorize/optimize/constfold.py @@ -6,21 +6,8 @@ import operator from typing import Any -from array_vectorize.ir import ( - Binding, - BinOp, - Call, - Compare, - DType, - Kind, - Literal, - Loop, - Node, - Program, - Stmt, - UnaryOp, -) -from array_vectorize.ir.walk import rewrite +from array_vectorize import ir +from array_vectorize.ir import walk __all__ = ["const_fold"] @@ -37,7 +24,7 @@ def _int64_ok(value: int) -> bool: return _INT64_MIN <= value <= _INT64_MAX -def _literal_kind(value: object) -> Kind: +def _literal_kind(value: object) -> ir.Kind: if isinstance(value, bool): return "bool" if isinstance(value, int): @@ -45,7 +32,7 @@ def _literal_kind(value: object) -> Kind: return "float" -def _fold_binop(op: str, left: Literal, right: Literal) -> Literal | None: +def _fold_binop(op: str, left: ir.Literal, right: ir.Literal) -> ir.Literal | None: if op in _NO_FOLD_OPS: return None if op in _ZERO_RISK_OPS and right.value == 0: @@ -82,14 +69,14 @@ def _fold_binop(op: str, left: Literal, right: Literal) -> Literal | None: return None if isinstance(result, int) and not isinstance(result, bool) and not _int64_ok(result): return None - return Literal(result, _literal_kind(result)) # type: ignore[arg-type] + return ir.Literal(result, _literal_kind(result)) # type: ignore[arg-type] -def _fold_unary(op: str, lit: Literal) -> Literal | None: +def _fold_unary(op: str, lit: ir.Literal) -> ir.Literal | None: v: Any = lit.value if op == "not": # exact per D2: logical_not(NaN) is False, Python `not nan` is False - return Literal(not v, "bool") + return ir.Literal(not v, "bool") try: match op: case "neg": @@ -104,7 +91,7 @@ def _fold_unary(op: str, lit: Literal) -> Literal | None: return None if isinstance(result, int) and not isinstance(result, bool) and not _int64_ok(result): return None - return Literal(result, _literal_kind(result)) + return ir.Literal(result, _literal_kind(result)) #: Array API functions that are pure and safe to fold over literal args, @@ -156,39 +143,39 @@ def _fold_unary(op: str, lit: Literal) -> Literal | None: } -def _fold_compare(op: str, left: Literal, right: Literal) -> Literal | None: +def _fold_compare(op: str, left: ir.Literal, right: ir.Literal) -> ir.Literal | None: try: result = _CMPOP_PY[op](left.value, right.value) except (ArithmeticError, ValueError, TypeError): return None - return Literal(bool(result), "bool") + return ir.Literal(bool(result), "bool") -def _const_fold_expr(node: Node) -> Node: - if isinstance(node, BinOp): - node = BinOp(node.op, _const_fold_expr(node.left), _const_fold_expr(node.right)) - if isinstance(node.left, Literal) and isinstance(node.right, Literal): +def _const_fold_expr(node: ir.Node) -> ir.Node: + if isinstance(node, ir.BinOp): + node = ir.BinOp(node.op, _const_fold_expr(node.left), _const_fold_expr(node.right)) + if isinstance(node.left, ir.Literal) and isinstance(node.right, ir.Literal): folded = _fold_binop(node.op, node.left, node.right) if folded is not None: return folded return node - if isinstance(node, UnaryOp): - node = UnaryOp(node.op, _const_fold_expr(node.operand)) - if isinstance(node.operand, Literal): + if isinstance(node, ir.UnaryOp): + node = ir.UnaryOp(node.op, _const_fold_expr(node.operand)) + if isinstance(node.operand, ir.Literal): folded = _fold_unary(node.op, node.operand) if folded is not None: return folded return node - if isinstance(node, Compare): - node = Compare(node.op, _const_fold_expr(node.left), _const_fold_expr(node.right)) - if isinstance(node.left, Literal) and isinstance(node.right, Literal): + if isinstance(node, ir.Compare): + node = ir.Compare(node.op, _const_fold_expr(node.left), _const_fold_expr(node.right)) + if isinstance(node.left, ir.Literal) and isinstance(node.right, ir.Literal): folded = _fold_compare(node.op, node.left, node.right) if folded is not None: return folded return node - node = rewrite(node, _const_fold_expr) - if isinstance(node, Call) and node.fn in _FOLDABLE_CALLS and node.args: - values = [a.value for a in node.args if isinstance(a, Literal)] + node = walk.rewrite(node, _const_fold_expr) + if isinstance(node, ir.Call) and node.fn in _FOLDABLE_CALLS and node.args: + values = [a.value for a in node.args if isinstance(a, ir.Literal)] if len(values) != len(node.args): return node try: @@ -197,41 +184,41 @@ def _const_fold_expr(node: Node) -> Node: pass else: if isinstance(value, bool): - return Literal(value, "bool") + return ir.Literal(value, "bool") if isinstance(value, int) and _int64_ok(value): - return Literal(value, "int") + return ir.Literal(value, "int") if isinstance(value, float): - return Literal(value, "float") + return ir.Literal(value, "float") # constant casts: int(3.14) -> 3 (also keeps plain-float astype out of # generated code, where xp.astype would fail on non-arrays) if ( - isinstance(node, Call) + isinstance(node, ir.Call) and node.fn == "astype" and len(node.args) == 2 - and isinstance(node.args[0], Literal) - and isinstance(node.args[1], DType) + and isinstance(node.args[0], ir.Literal) + and isinstance(node.args[1], ir.DType) ): value, dtype = node.args[0].value, node.args[1].name try: if dtype == "bool": - return Literal(bool(value), "bool") + return ir.Literal(bool(value), "bool") if dtype == "float64": - return Literal(float(value), "float") + return ir.Literal(float(value), "float") if _int64_ok(int(value)): - return Literal(int(value), "int") + return ir.Literal(int(value), "int") except (ValueError, OverflowError): pass # int(inf)/int(nan) cannot fold (scalar Python raises Overflow/ # ValueError): keep the cast but wrap the literal so xp.astype gets # an array, not a plain float (backend-defined result, design D3) - return Call("astype", (Call("asarray", (node.args[0],)), node.args[1])) + return ir.Call("astype", (ir.Call("asarray", (node.args[0],)), node.args[1])) return node -def _fold_stmt(stmt: Stmt) -> Stmt: - if isinstance(stmt, Binding): - return Binding(stmt.name, _const_fold_expr(stmt.expr)) - return Loop( +def _fold_stmt(stmt: ir.Stmt) -> ir.Stmt: + if isinstance(stmt, ir.Binding): + return ir.Binding(stmt.name, _const_fold_expr(stmt.expr)) + return ir.Loop( stmt.var, _const_fold_expr(stmt.start), _const_fold_expr(stmt.stop), @@ -240,7 +227,7 @@ def _fold_stmt(stmt: Stmt) -> Stmt: ) -def const_fold(program: Program) -> Program: +def const_fold(program: ir.Program) -> ir.Program: """Constant-fold every binding and the result expression.""" bindings = tuple(_fold_stmt(s) for s in program.bindings) - return Program(program.params, bindings, _const_fold_expr(program.result)) + return ir.Program(program.params, bindings, _const_fold_expr(program.result)) diff --git a/src/array_vectorize/optimize/cse.py b/src/array_vectorize/optimize/cse.py index 42341ce..0700b1e 100644 --- a/src/array_vectorize/optimize/cse.py +++ b/src/array_vectorize/optimize/cse.py @@ -2,36 +2,35 @@ from __future__ import annotations -from collections import Counter +import collections -from array_vectorize.ir import Binding, Call, Loop, Node, Program, Ref, Stmt, Where, walk -from array_vectorize.ir.ssa import SSAEnv -from array_vectorize.ir.walk import children +from array_vectorize import ir +from array_vectorize.ir import ssa, walk __all__ = ["cse"] # ----------------------------------------------------------------------- CSE -_CSE_ELIGIBLE = (Call, Where) +_CSE_ELIGIBLE = (ir.Call, ir.Where) -def _count_eligible(node: Node, counter: Counter[Node]) -> None: +def _count_eligible(node: ir.Node, counter: collections.Counter[ir.Node]) -> None: if isinstance(node, _CSE_ELIGIBLE): counter[node] += 1 - for child in children(node): + for child in walk.children(node): _count_eligible(child, counter) -def _has_loop(stmts: tuple[Stmt, ...] | list[Stmt]) -> bool: +def _has_loop(stmts: tuple[ir.Stmt, ...] | list[ir.Stmt]) -> bool: for stmt in stmts: - if isinstance(stmt, Loop): + if isinstance(stmt, ir.Loop): if _has_loop(stmt.body): return True return True return False -def cse(program: Program, ssa: SSAEnv) -> Program: +def cse(program: ir.Program, ssa: ssa.SSAEnv) -> ir.Program: """Hoist duplicated Call/Where subtrees into temps ``t_1, t_2, ...``. Each pass counts eligible nodes in the *current* program and hoists those @@ -48,42 +47,42 @@ def cse(program: Program, ssa: SSAEnv) -> Program: # needed for correctness). return current for _ in range(8): - counter: Counter[Node] = Counter() + counter: collections.Counter[ir.Node] = collections.Counter() for stmt in current.bindings: - assert isinstance(stmt, Binding) # CSE skips loop programs + assert isinstance(stmt, ir.Binding) # CSE skips loop programs _count_eligible(stmt.expr, counter) _count_eligible(current.result, counter) if not any(count >= 2 for count in counter.values()): return current - memo: dict[Node, str] = {} - pending: list[Binding] = [] + memo: dict[ir.Node, str] = {} + pending: list[ir.Binding] = [] def rewrite( - node: Node, + node: ir.Node, *, - counter: Counter[Node] = counter, - memo: dict[Node, str] = memo, - pending: list[Binding] = pending, - ) -> Node: + counter: collections.Counter[ir.Node] = counter, + memo: dict[ir.Node, str] = memo, + pending: list[ir.Binding] = pending, + ) -> ir.Node: node = walk.rewrite(node, rewrite) if isinstance(node, _CSE_ELIGIBLE) and counter[node] >= 2: if node not in memo: temp = ssa.fresh_temp("t") - pending.append(Binding(temp, node)) + pending.append(ir.Binding(temp, node)) memo[node] = temp - return Ref(memo[node]) + return ir.Ref(memo[node]) return node - rewritten: list[Stmt] = [] + rewritten: list[ir.Stmt] = [] for stmt in current.bindings: - assert isinstance(stmt, Binding) # CSE skips loop programs + assert isinstance(stmt, ir.Binding) # CSE skips loop programs start = len(pending) expr = rewrite(stmt.expr) rewritten.extend(pending[start:]) - rewritten.append(Binding(stmt.name, expr)) + rewritten.append(ir.Binding(stmt.name, expr)) result_start = len(pending) result = rewrite(current.result) rewritten.extend(pending[result_start:]) # temps first used by the result go last - current = Program(current.params, tuple(rewritten), result) + current = ir.Program(current.params, tuple(rewritten), result) return current diff --git a/src/array_vectorize/optimize/dce.py b/src/array_vectorize/optimize/dce.py index ba731e9..5fabec6 100644 --- a/src/array_vectorize/optimize/dce.py +++ b/src/array_vectorize/optimize/dce.py @@ -2,24 +2,24 @@ from __future__ import annotations -from array_vectorize.ir import Binding, Loop, Node, Program, Ref, Stmt -from array_vectorize.ir.walk import children +from array_vectorize import ir +from array_vectorize.ir import walk __all__ = ["dce"] # ----------------------------------------------------------------------- DCE -def _mark_live(node: Node, live: set[str]) -> None: - if isinstance(node, Ref): +def _mark_live(node: ir.Node, live: set[str]) -> None: + if isinstance(node, ir.Ref): live.add(node.name) - for child in children(node): + for child in walk.children(node): _mark_live(child, live) -def _mark_live_stmt(stmt: Stmt, live: set[str]) -> Stmt | None: +def _mark_live_stmt(stmt: ir.Stmt, live: set[str]) -> ir.Stmt | None: """Mark liveness from a statement; returns the kept statement or None.""" - if isinstance(stmt, Binding): + if isinstance(stmt, ir.Binding): if stmt.name in live: _mark_live(stmt.expr, live) return stmt @@ -28,10 +28,10 @@ def _mark_live_stmt(stmt: Stmt, live: set[str]) -> Stmt | None: # fixed point: a statement can be live only through a use appearing # EARLIER in the body (loop-carried feedback), which a single reverse # pass misses. Liveness only grows, so this terminates. - kept: list[Stmt] = [] + kept: list[ir.Stmt] = [] while True: before_live = set(live) - new_kept: list[Stmt] = [] + new_kept: list[ir.Stmt] = [] for inner in reversed(stmt.body): kept_inner = _mark_live_stmt(inner, live) if kept_inner is not None: @@ -45,19 +45,19 @@ def _mark_live_stmt(stmt: Stmt, live: set[str]) -> Stmt | None: _mark_live(stmt.stop, live) _mark_live(stmt.step, live) if kept_any or stmt.var in live: - return Loop(stmt.var, stmt.start, stmt.stop, stmt.step, tuple(kept)) + return ir.Loop(stmt.var, stmt.start, stmt.stop, stmt.step, tuple(kept)) return None -def dce(program: Program) -> Program: +def dce(program: ir.Program) -> ir.Program: """Drop bindings whose results are never used.""" live: set[str] = set() _mark_live(program.result, live) - kept: list[Stmt] = [] + kept: list[ir.Stmt] = [] # single reverse pass suffices: SSA uses only earlier bindings for stmt in reversed(program.bindings): kept_stmt = _mark_live_stmt(stmt, live) if kept_stmt is not None: kept.append(kept_stmt) kept.reverse() - return Program(program.params, tuple(kept), program.result) + return ir.Program(program.params, tuple(kept), program.result) diff --git a/src/array_vectorize/optimize/protect_domains.py b/src/array_vectorize/optimize/protect_domains.py index 4574375..db10220 100644 --- a/src/array_vectorize/optimize/protect_domains.py +++ b/src/array_vectorize/optimize/protect_domains.py @@ -2,22 +2,9 @@ from __future__ import annotations -from array_vectorize.ir import ( - Binding, - Call, - Compare, - Literal, - Logical, - Loop, - Node, - Program, - Ref, - Stmt, - UnaryOp, - Where, - walk, -) -from array_vectorize.ir.walk import children +from array_vectorize import ir +from array_vectorize.ir import walk +from array_vectorize.ir import walk as walk_mod __all__ = ["protect_domains"] @@ -40,19 +27,19 @@ } -def _and_ctx(ctx: Node | None, cond: Node) -> Node: - return cond if ctx is None else Logical("and", (ctx, cond)) +def _and_ctx(ctx: ir.Node | None, cond: ir.Node) -> ir.Node: + return cond if ctx is None else ir.Logical("and", (ctx, cond)) -def _free_names(node: Node) -> set[str]: +def _free_names(node: ir.Node) -> set[str]: """All Ref names appearing in a node tree.""" names: set[str] = set() - def walk(n: Node) -> None: - if isinstance(n, Ref): + def walk(n: ir.Node) -> None: + if isinstance(n, ir.Ref): names.add(n.name) return - for child in children(n): + for child in walk_mod.children(n): walk(child) walk(node) @@ -60,31 +47,31 @@ def walk(n: Node) -> None: def _clamp_dead( - arg: Node, - dead: Node, + arg: ir.Node, + dead: ir.Node, spec: tuple[float, bool, float | None, bool, float, float | None], -) -> Node: +) -> ir.Node: """Clamp ``arg`` to the domain, but only where ``dead`` holds.""" lo, lo_open, hi, hi_open, safe_lo, safe_hi = spec - clamped: Node = arg + clamped: ir.Node = arg if hi is not None: assert safe_hi is not None op = "ge" if hi_open else "gt" - clamped = Where( - Logical("and", (dead, Compare(op, arg, Literal(hi, "float")))), - Literal(safe_hi, "float"), + clamped = ir.Where( + ir.Logical("and", (dead, ir.Compare(op, arg, ir.Literal(hi, "float")))), + ir.Literal(safe_hi, "float"), clamped, ) op = "le" if lo_open else "lt" - clamped = Where( - Logical("and", (dead, Compare(op, arg, Literal(lo, "float")))), - Literal(safe_lo, "float"), + clamped = ir.Where( + ir.Logical("and", (dead, ir.Compare(op, arg, ir.Literal(lo, "float")))), + ir.Literal(safe_lo, "float"), clamped, ) return clamped -def protect_domains(program: Program) -> Program: +def protect_domains(program: ir.Program) -> ir.Program: """Clamp partial-function arguments on dead lanes only (design D1). A single reverse pass propagates liveness conditions: each binding's @@ -94,9 +81,9 @@ def protect_domains(program: Program) -> Program: never change and dead lanes never warn. Loop bodies are left unprotected (conservative) and their reads mark bindings fully live. """ - live: dict[str, Node | None] = {} + live: dict[str, ir.Node | None] = {} - def or_live(name: str, ctx: Node | None) -> None: + def or_live(name: str, ctx: ir.Node | None) -> None: if name not in live: live[name] = ctx return @@ -104,13 +91,13 @@ def or_live(name: str, ctx: Node | None) -> None: if cur is None or ctx is None: live[name] = None else: - live[name] = Logical("or", (cur, ctx)) + live[name] = ir.Logical("or", (cur, ctx)) - def record_uses(node: Node, ctx: Node | None, skip: frozenset[str] = frozenset()) -> None: - if isinstance(node, Ref): + def record_uses(node: ir.Node, ctx: ir.Node | None, skip: frozenset[str] = frozenset()) -> None: + if isinstance(node, ir.Ref): if node.name not in skip: or_live(node.name, ctx) - elif isinstance(node, Where): + elif isinstance(node, ir.Where): record_uses(node.cond, ctx, skip) # the cond evaluates on every lane, so every name it reads is # already live under the full ctx; re-recording those names @@ -118,42 +105,42 @@ def record_uses(node: Node, ctx: Node | None, skip: frozenset[str] = frozenset() # and possibly self-referential — disjuncts to their liveness cond_names = frozenset(n for n in _free_names(node.cond) if n not in skip) record_uses(node.then, _and_ctx(ctx, node.cond), skip | cond_names) - record_uses(node.other, _and_ctx(ctx, UnaryOp("not", node.cond)), skip | cond_names) + record_uses(node.other, _and_ctx(ctx, ir.UnaryOp("not", node.cond)), skip | cond_names) else: - for child in children(node): + for child in walk_mod.children(node): record_uses(child, ctx, skip) - def rewrite(node: Node, ctx: Node | None) -> Node: - if isinstance(node, Ref): + def rewrite(node: ir.Node, ctx: ir.Node | None) -> ir.Node: + if isinstance(node, ir.Ref): return node - if isinstance(node, Where): + if isinstance(node, ir.Where): cond = rewrite(node.cond, ctx) then = rewrite(node.then, _and_ctx(ctx, cond)) - other = rewrite(node.other, _and_ctx(ctx, UnaryOp("not", cond))) - return Where(cond, then, other) - if isinstance(node, Call) and node.fn in _PARTIAL_DOMAINS: + other = rewrite(node.other, _and_ctx(ctx, ir.UnaryOp("not", cond))) + return ir.Where(cond, then, other) + if isinstance(node, ir.Call) and node.fn in _PARTIAL_DOMAINS: args = [rewrite(a, ctx) for a in node.args] if ctx is not None: - dead = UnaryOp("not", ctx) + dead = ir.UnaryOp("not", ctx) spec = _PARTIAL_DOMAINS[node.fn] args = [_clamp_dead(a, dead, spec) for a in args] - return Call(node.fn, tuple(args)) + return ir.Call(node.fn, tuple(args)) return walk.rewrite(node, lambda n: rewrite(n, ctx)) - def mark_fully_live(node: Node) -> None: - if isinstance(node, Ref): + def mark_fully_live(node: ir.Node) -> None: + if isinstance(node, ir.Ref): live[node.name] = None else: - for child in children(node): + for child in walk_mod.children(node): mark_fully_live(child) - def rewrite_stmt(stmt: Stmt, ahead: set[str]) -> Stmt: + def rewrite_stmt(stmt: ir.Stmt, ahead: set[str]) -> ir.Stmt: """Rewrite one binding. ``ahead`` holds the names bound AFTER this statement (they are not yet assigned where this statement runs). """ - if isinstance(stmt, Binding): + if isinstance(stmt, ir.Binding): ctx = live.pop(stmt.name, None) record_ctx = ctx if ctx is not None and _free_names(ctx) & (ahead | {stmt.name}): @@ -165,23 +152,23 @@ def rewrite_stmt(stmt: Stmt, ahead: set[str]) -> Stmt: expr = rewrite(stmt.expr, ctx) record_uses(expr, record_ctx) ahead.add(stmt.name) - return Binding(stmt.name, expr) - assert isinstance(stmt, Loop) + return ir.Binding(stmt.name, expr) + assert isinstance(stmt, ir.Loop) # conservative: loop bodies are not rewritten; every name they read # is marked fully live (later iterations may consume any binding) mark_fully_live(stmt.start) mark_fully_live(stmt.stop) mark_fully_live(stmt.step) for inner in stmt.body: - if isinstance(inner, Binding): + if isinstance(inner, ir.Binding): mark_fully_live(inner.expr) # everything bound inside the loop (recursively) is not yet # assigned where statements BEFORE the loop run — loop-body # bindings count as forward references just like the index ahead.add(stmt.var) - def collect_bound(st: Stmt) -> None: - if isinstance(st, Binding): + def collect_bound(st: ir.Stmt) -> None: + if isinstance(st, ir.Binding): ahead.add(st.name) else: ahead.add(st.var) @@ -197,4 +184,4 @@ def collect_bound(st: Stmt) -> None: ahead: set[str] = set() rewritten = [rewrite_stmt(stmt, ahead) for stmt in reversed(program.bindings)] rewritten.reverse() - return Program(program.params, tuple(rewritten), new_result) + return ir.Program(program.params, tuple(rewritten), new_result) diff --git a/src/array_vectorize/pipeline.py b/src/array_vectorize/pipeline.py index b19fec2..3a3df14 100644 --- a/src/array_vectorize/pipeline.py +++ b/src/array_vectorize/pipeline.py @@ -12,13 +12,14 @@ from collections.abc import Callable from typing import Any -from array_vectorize.codegen import generate_source -from array_vectorize.emit import compile_vectorized -from array_vectorize.frontend.extract import extract_function, resolve_original -from array_vectorize.frontend.validate import validate -from array_vectorize.lower import HelperVectorizer, lower_function -from array_vectorize.optimize import optimize, protect_domains -from array_vectorize.verify import verify_match +from array_vectorize import codegen, emit, lower, verify +from array_vectorize import optimize as optimize_mod +from array_vectorize.frontend import extract +from array_vectorize.frontend import validate as validate_mod + +# optimize's documented submodule/function name collision (the NOTE in +# optimize/__init__.py): the from-import is the supported spelling +from array_vectorize.optimize import protect_domains # cleanporter: ignore[CP002] name collision __all__ = ["compile_function"] @@ -28,7 +29,7 @@ def compile_function( *, protect: bool = False, verify_args: tuple[Any, ...] | None = None, - helper_vectorizer: HelperVectorizer | None = None, + helper_vectorizer: lower.HelperVectorizer | None = None, namespace: Any = None, ) -> Callable[..., Any]: """Run the strict pipeline over ``func``. @@ -53,24 +54,26 @@ def compile_function( (so all-scalar calls become legal); the namespace is threaded to helper vectorization and verification. """ - info = extract_function(func) - validate(info) - lowered = lower_function(info, helper_vectorizer=helper_vectorizer) - program = optimize(lowered.program, user_names=info.user_names | lowered.emitted_names) + info = extract.extract_function(func) + validate_mod.validate(info) + lowered = lower.lower_function(info, helper_vectorizer=helper_vectorizer) + program = optimize_mod.optimize( + lowered.program, user_names=info.user_names | lowered.emitted_names + ) if protect: program = protect_domains(program) - source = generate_source(lowered, program, pinned=namespace is not None) + source = codegen.generate_source(lowered, program, pinned=namespace is not None) hidden_params = lowered.hidden_params if namespace is not None: # the pinned namespace rides as the hidden kw-only parameter's # runtime default (the existing __kwdefaults__ injection handles it) hidden_params = [*lowered.hidden_params, (lowered.namespace_param, namespace)] - vec = compile_vectorized( + vec = emit.compile_vectorized( source, lowered.name, hidden_params, original=func, helpers=lowered.helpers ) if verify_args is not None: # compare against the resolved scalar original (what was actually # compiled), not the callable as passed — a prior vectorization # passed back in cannot serve as the element-wise oracle - verify_match(vec, resolve_original(func), verify_args, namespace=namespace) + verify.verify_match(vec, extract.resolve_original(func), verify_args, namespace=namespace) return vec diff --git a/src/array_vectorize/runtime/arith.py b/src/array_vectorize/runtime/arith.py index 4971144..b6dcadf 100644 --- a/src/array_vectorize/runtime/arith.py +++ b/src/array_vectorize/runtime/arith.py @@ -10,7 +10,7 @@ import operator from typing import Any -from array_vectorize.runtime.dtype import _describe_dtype, _dtype_bits +from array_vectorize.runtime import dtype __all__ = [ "_ARITH_FNS", @@ -57,7 +57,7 @@ class ArithOp(enum.IntEnum): def _is_u64(a: Any) -> bool: - return hasattr(a, "dtype") and "uint" in str(a.dtype) and _dtype_bits(a.dtype) >= 64 + return hasattr(a, "dtype") and "uint" in str(a.dtype) and dtype._dtype_bits(a.dtype) >= 64 def _to_u64(xp: Any, a: Any) -> Any: @@ -226,7 +226,7 @@ def _vec_arith(xp: Any, op_id: Any, left: Any, right: Any = None) -> Any: xp.asarray(left, dtype=xp.float64), xp.asarray(right, dtype=xp.float64) ) # generic: cast to the common dtype and compute - classes = [_describe_dtype(a) for a in args] + classes = [dtype._describe_dtype(a) for a in args] has_float = any(f for f, _ in classes) has_int = any(not f for f, _ in classes) if op == "div" or (has_float and has_int): diff --git a/src/array_vectorize/runtime/minmax.py b/src/array_vectorize/runtime/minmax.py index 54ac89a..246aa66 100644 --- a/src/array_vectorize/runtime/minmax.py +++ b/src/array_vectorize/runtime/minmax.py @@ -8,12 +8,7 @@ from typing import Any -from array_vectorize.runtime.dtype import ( - _describe_dtype, - _dtype_bits, - _fits_dtype, - _minmax_bound_kind, -) +from array_vectorize.runtime import dtype __all__ = ["_vec_minmax"] @@ -82,18 +77,18 @@ def fold(cast: list[Any]) -> Any: if arrays: first_dt = arrays[0].dtype if all(a.dtype == first_dt for a in arrays): - if all(_fits_dtype(b, first_dt) for b in lits): + if all(dtype._fits_dtype(b, first_dt) for b in lits): cast = [a if hasattr(a, "dtype") else xp.asarray(a, dtype=first_dt) for a in args] return fold(cast) - if all(_minmax_bound_kind(b, is_min, first_dt) != "promote" for b in lits): + if all(dtype._minmax_bound_kind(b, is_min, first_dt) != "promote" for b in lits): cast = [] for a in args: if hasattr(a, "dtype"): cast.append(a) - elif _fits_dtype(a, first_dt): + elif dtype._fits_dtype(a, first_dt): cast.append(xp.asarray(a, dtype=first_dt)) else: # clamp: the bound can never win - bound = 2 ** _dtype_bits(first_dt) - 1 if is_min else 0 + bound = 2 ** dtype._dtype_bits(first_dt) - 1 if is_min else 0 cast.append(xp.asarray(bound, dtype=first_dt)) return fold(cast) @@ -101,7 +96,7 @@ def fold(cast: list[Any]) -> Any: isinstance(b, float) for b in lits ) if has_float: - classes = [_describe_dtype(a) for a in args] + classes = [dtype._describe_dtype(a) for a in args] if any(not f for f, _ in classes): dt = xp.float64 else: @@ -109,7 +104,7 @@ def fold(cast: list[Any]) -> Any: return fold([xp.asarray(a, dtype=dt) for a in args]) # pure-integer domain - u64_present = any("uint" in str(a.dtype) and _dtype_bits(a.dtype) >= 64 for a in arrays) + u64_present = any("uint" in str(a.dtype) and dtype._dtype_bits(a.dtype) >= 64 for a in arrays) negative_capable = any( "int" in str(a.dtype) and "uint" not in str(a.dtype) for a in arrays ) or any(isinstance(b, int) and not isinstance(b, bool) and b < 0 for b in lits) @@ -119,7 +114,7 @@ def fold(cast: list[Any]) -> Any: cast = [] for a in args: if hasattr(a, "dtype"): - if "uint" in str(a.dtype) and _dtype_bits(a.dtype) >= 64: + if "uint" in str(a.dtype) and dtype._dtype_bits(a.dtype) >= 64: cast.append( xp.astype( xp.where(a > imax, xp.asarray(imax, dtype=a.dtype), a), @@ -145,6 +140,6 @@ def fold(cast: list[Any]) -> Any: else: cast.append(xp.asarray(a, dtype=xp.uint64)) return fold(cast) - classes = [_describe_dtype(a) for a in args] + classes = [dtype._describe_dtype(a) for a in args] dt = getattr(xp, f"int{max(w for _, w in classes)}") return fold([xp.asarray(a, dtype=dt) for a in args]) diff --git a/src/array_vectorize/runtime/registry.py b/src/array_vectorize/runtime/registry.py index be4e9cf..f443c89 100644 --- a/src/array_vectorize/runtime/registry.py +++ b/src/array_vectorize/runtime/registry.py @@ -11,8 +11,11 @@ from collections.abc import Callable from typing import Any -from array_vectorize.runtime.arith import ArithOp, _vec_arith -from array_vectorize.runtime.minmax import _vec_minmax +from array_vectorize.runtime import arith, minmax + +# lowering goes through the registry seam (docs/architecture.md): +# lower/expressions.py imports ArithOp from here, not from runtime.arith +from array_vectorize.runtime.arith import ArithOp # cleanporter: ignore[CP003] re-export __all__ = ["RUNTIME_HELPERS", "ArithOp"] @@ -20,6 +23,6 @@ #: The base names are load-bearing: generated source (and the golden #: snapshots) reference them verbatim. RUNTIME_HELPERS: dict[str, tuple[str, Callable[..., Any]]] = { - "vec_arith": ("_vec_arith", _vec_arith), - "vec_minmax": ("_vec_minmax", _vec_minmax), + "vec_arith": ("_vec_arith", arith._vec_arith), + "vec_minmax": ("_vec_minmax", minmax._vec_minmax), } diff --git a/src/array_vectorize/verify.py b/src/array_vectorize/verify.py index 07a9df3..7235f71 100644 --- a/src/array_vectorize/verify.py +++ b/src/array_vectorize/verify.py @@ -12,10 +12,9 @@ from collections.abc import Callable from typing import Any -from array_api_compat import array_namespace +import array_api_compat -from array_vectorize.compat import _is_array, _py_scalar -from array_vectorize.errors import VectorizationError +from array_vectorize import compat, errors __all__ = ["verify_match"] @@ -34,11 +33,13 @@ def verify_match( With ``namespace`` set (a pinned vectorization), that namespace is used directly instead of extracting it from the example arrays. """ - positions = [i for i, a in enumerate(example_args) if _is_array(a)] + positions = [i for i, a in enumerate(example_args) if compat._is_array(a)] if not positions: - raise VectorizationError("verify= needs at least one Array API array in the example inputs") + raise errors.VectorizationError( + "verify= needs at least one Array API array in the example inputs" + ) arrays = [example_args[i] for i in positions] - xp = namespace if namespace is not None else array_namespace(*arrays) + xp = namespace if namespace is not None else array_api_compat.array_namespace(*arrays) with warnings.catch_warnings(): # out-of-domain lanes are expected (design D3); their warnings are noise # (including inf-inf in the comparison below) @@ -55,7 +56,7 @@ def verify_match( if got_arr.shape == (): got_arr = xp.broadcast_to(got_arr, shape) if tuple(got_arr.shape) != tuple(shape): - raise VectorizationError( + raise errors.VectorizationError( f"verification failed: result shape {tuple(got_arr.shape)} " f"!= input shape {tuple(shape)}" ) @@ -69,7 +70,7 @@ def verify_match( # int/bool arithmetic semantics) call_args = list(example_args) for pos, elem in zip(positions, elems, strict=True): - call_args[pos] = _py_scalar(elem) + call_args[pos] = compat._py_scalar(elem) try: expected_values.append(original(*call_args)) except (ValueError, ArithmeticError, OverflowError): @@ -92,7 +93,7 @@ def verify_match( close = (got_f == exp_f) | (xp.isfinite(diff) & (diff <= _ATOL + _RTOL * xp.abs(exp_f))) both_nan = xp.isnan(got_f) & xp.isnan(exp_f) if not bool(xp.all(close | both_nan)): - raise VectorizationError( + raise errors.VectorizationError( "verification failed: vectorized output differs from the scalar " f"original on the provided example inputs (got {got}, expected {expected})" ) diff --git a/tests/features/test_namespace_pin.py b/tests/features/test_namespace_pin.py index f43eea8..9c94545 100644 --- a/tests/features/test_namespace_pin.py +++ b/tests/features/test_namespace_pin.py @@ -8,16 +8,16 @@ from __future__ import annotations import inspect +import pathlib import warnings -from pathlib import Path import corpus import numpy as np import pytest -from support import make_fn, make_module +import support -from array_vectorize import vectorize -from array_vectorize.api import _HELPER_CACHE +import array_vectorize +from array_vectorize import api X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) Y = np.asarray([3.0, 1.5, 1.0, 0.5, -1.0, -3.0]) @@ -27,7 +27,7 @@ def test_pinned_source_shape() -> None: - vec = vectorize(corpus.add, namespace=np) + vec = array_vectorize.vectorize(corpus.add, namespace=np) assert "def add_vec(x, y, *, _namespace=None):" in vec.source assert "xp = _namespace" in vec.source # pinned generated source has zero imports and no namespace extraction @@ -37,10 +37,10 @@ def test_pinned_source_shape() -> None: def test_unpinned_source_unchanged() -> None: - unpinned = vectorize(corpus.add) - pinned = vectorize(corpus.add, namespace=np) + unpinned = array_vectorize.vectorize(corpus.add) + pinned = array_vectorize.vectorize(corpus.add, namespace=np) # byte-identical to the pre-pin generated source (the golden snapshot) - golden = (Path(__file__).parent.parent / "golden" / "cases" / "add.py").read_text() + golden = (pathlib.Path(__file__).parent.parent / "golden" / "cases" / "add.py").read_text() assert unpinned.source == golden assert pinned.source != unpinned.source assert "xp = _namespace" not in unpinned.source @@ -48,7 +48,7 @@ def test_unpinned_source_unchanged() -> None: def test_pinned_closure_array_still_hidden_param() -> None: - vec = vectorize(corpus.closure_array, namespace=np) + vec = array_vectorize.vectorize(corpus.closure_array, namespace=np) assert "*, ARR=None, _namespace=None" in vec.source assert vec.__kwdefaults__["_namespace"] is np assert np.allclose(vec(np.asarray([0.0, 1.0, 2.0])), [1.0, 3.0, 5.0]) @@ -61,15 +61,19 @@ def test_pinned_closure_array_still_hidden_param() -> None: def test_pinned_numpy_matches_unpinned(name: str) -> None: fn = getattr(corpus, name) args = (X, Y) if name == "add" else (X,) - assert np.allclose(vectorize(fn, namespace=np)(*args), vectorize(fn)(*args), equal_nan=True) + assert np.allclose( + array_vectorize.vectorize(fn, namespace=np)(*args), + array_vectorize.vectorize(fn)(*args), + equal_nan=True, + ) def test_pinned_strict_matches_unpinned() -> None: import array_api_strict as xps a = xps.asarray([1.0, -2.0, 3.0]) - unpinned = vectorize(corpus.relu) - pinned = vectorize(corpus.relu, namespace=xps) + unpinned = array_vectorize.vectorize(corpus.relu) + pinned = array_vectorize.vectorize(corpus.relu, namespace=xps) assert list(map(float, pinned(a))) == list(map(float, unpinned(a))) @@ -77,20 +81,20 @@ def test_pinned_strict_matches_unpinned() -> None: def test_all_scalar_call_with_pin() -> None: - fn = make_fn(" return math.sqrt(x * x)", defaults="x") - vec = vectorize(fn, namespace=np) + fn = support.make_fn(" return math.sqrt(x * x)", defaults="x") + vec = array_vectorize.vectorize(fn, namespace=np) assert float(vec(3.0)) == pytest.approx(3.0) def test_all_scalar_minmax_with_pin() -> None: - fn = make_fn(" return min(x, 3.0) + math.sqrt(x)", defaults="x") - vec = vectorize(fn, namespace=np) + fn = support.make_fn(" return min(x, 3.0) + math.sqrt(x)", defaults="x") + vec = array_vectorize.vectorize(fn, namespace=np) assert float(vec(4.0)) == pytest.approx(3.0 + 2.0) def test_all_scalar_call_without_pin_raises() -> None: - fn = make_fn(" return math.sqrt(x * x)", defaults="x") - vec = vectorize(fn) + fn = support.make_fn(" return math.sqrt(x * x)", defaults="x") + vec = array_vectorize.vectorize(fn) # message wording varies across array-api-compat versions: 1.10 says # "Unrecognized array input", newer releases "requires at least one # non-scalar array input" — the contract is the TypeError itself @@ -102,26 +106,26 @@ def test_all_scalar_call_without_pin_raises() -> None: def test_with_namespace_equivalent_results() -> None: - vf = vectorize(corpus.psi) + vf = array_vectorize.vectorize(corpus.psi) assert np.allclose(vf.with_namespace(np)(X), vf(X)) def test_with_namespace_identity_stable() -> None: - vf = vectorize(corpus.psi) + vf = array_vectorize.vectorize(corpus.psi) assert vf.with_namespace(np) is vf.with_namespace(np) def test_with_namespace_memoized_across_compilations() -> None: # the cache keys on the scalar original, not on the compiled wrapper: # distinct compilations (protected ones are never memoized) share pins - vf1 = vectorize(corpus.psi, protect_domains=True) - vf2 = vectorize(corpus.psi, protect_domains=True) + vf1 = array_vectorize.vectorize(corpus.psi, protect_domains=True) + vf2 = array_vectorize.vectorize(corpus.psi, protect_domains=True) assert vf1 is not vf2 assert vf1.with_namespace(np) is vf2.with_namespace(np) def test_with_namespace_does_not_mutate() -> None: - vf = vectorize(corpus.psi) + vf = array_vectorize.vectorize(corpus.psi) w = vf.with_namespace(np) assert w is not vf # the original stays unpinned: no hidden namespace parameter at all @@ -132,29 +136,31 @@ def test_with_namespace_does_not_mutate() -> None: def test_with_namespace_chaining_replaces_pin() -> None: import array_api_strict as xps - vf = vectorize(corpus.psi) + vf = array_vectorize.vectorize(corpus.psi) a = xps.asarray([-2.0, 0.5, 3.0]) chained = vf.with_namespace(np).with_namespace(xps) assert list(map(float, chained(a))) == list(map(float, vf(a))) def test_with_namespace_preserves_contract() -> None: - vf = vectorize(corpus.psi) + vf = array_vectorize.vectorize(corpus.psi) w = vf.with_namespace(np) assert w._vectorized_original is corpus.psi # the true scalar original assert inspect.signature(w) == inspect.signature(corpus.psi) - assert w.source == vectorize(corpus.psi, namespace=np).source + assert w.source == array_vectorize.vectorize(corpus.psi, namespace=np).source assert w.__kwdefaults__["_namespace"] is np assert "psi_vec" in inspect.getsource(w) # linecache registration intact def test_with_namespace_does_not_rerun_verify(monkeypatch: pytest.MonkeyPatch) -> None: - import array_vectorize.pipeline as pipeline_mod + # pipeline calls verify.verify_match at call time, so patching the + # verify module's attribute is what the compiled path sees + import array_vectorize.verify as verify_mod calls: list[int] = [] - monkeypatch.setattr(pipeline_mod, "verify_match", lambda *a, **k: calls.append(1)) - fn = make_fn(" return x + 1.0", defaults="x") - vec = vectorize(fn, verify=(np.asarray([1.0]),)) + monkeypatch.setattr(verify_mod, "verify_match", lambda *a, **k: calls.append(1)) + fn = support.make_fn(" return x + 1.0", defaults="x") + vec = array_vectorize.vectorize(fn, verify=(np.asarray([1.0]),)) assert calls # the initial decoration-time verify ran calls.clear() vec.with_namespace(np) @@ -165,7 +171,7 @@ def test_with_namespace_does_not_rerun_verify(monkeypatch: pytest.MonkeyPatch) - def test_pinned_helper_all_scalar_call() -> None: - mod = make_module( + mod = support.make_module( "def helper_inner(y):\n" " return y * y + 1.0\n" "\n" @@ -173,14 +179,14 @@ def test_pinned_helper_all_scalar_call() -> None: "def helper_outer(x):\n" " return helper_inner(x) + helper_inner(x * 2.0)\n" ) - vec = vectorize(mod.helper_outer, namespace=np) + vec = array_vectorize.vectorize(mod.helper_outer, namespace=np) # only legal because the helper is pinned too (its own array_namespace # extraction would fail on scalars) assert float(vec(2.0)) == (2.0**2 + 1.0) + (4.0**2 + 1.0) def test_pinned_and_unpinned_helpers_are_distinct() -> None: - mod = make_module( + mod = support.make_module( "def helper_inner(y):\n" " return y * y + 1.0\n" "\n" @@ -188,10 +194,10 @@ def test_pinned_and_unpinned_helpers_are_distinct() -> None: "def helper_outer(x):\n" " return helper_inner(x) + helper_inner(x * 2.0)\n" ) - vectorize(mod.helper_outer) - vectorize(mod.helper_outer, namespace=np) - unpinned = _HELPER_CACHE[(mod.helper_inner, False, None)] - pinned = _HELPER_CACHE[(mod.helper_inner, False, np)] + array_vectorize.vectorize(mod.helper_outer) + array_vectorize.vectorize(mod.helper_outer, namespace=np) + unpinned = api._HELPER_CACHE[(mod.helper_inner, False, None)] + pinned = api._HELPER_CACHE[(mod.helper_inner, False, np)] assert unpinned is not pinned assert "xp = _namespace" in pinned.source assert "array_namespace" in unpinned.source @@ -201,8 +207,8 @@ def test_pinned_and_unpinned_helpers_are_distinct() -> None: def test_user_param_named_namespace_is_kept() -> None: - fn = make_fn(" return x + _namespace", defaults="x, _namespace") - vec = vectorize(fn, namespace=np) + fn = support.make_fn(" return x + _namespace", defaults="x, _namespace") + vec = array_vectorize.vectorize(fn, namespace=np) # the user's parameter keeps its name; ours is mangled assert "def subject_vec(x, _namespace, *, _namespace_1=None):" in vec.source assert "xp = _namespace_1" in vec.source @@ -213,8 +219,8 @@ def test_user_param_named_namespace_is_kept() -> None: def test_user_local_named_namespace_is_kept() -> None: - fn = make_fn(" _namespace = x + 1\n return _namespace * 2", defaults="x") - vec = vectorize(fn, namespace=np) + fn = support.make_fn(" _namespace = x + 1\n return _namespace * 2", defaults="x") + vec = array_vectorize.vectorize(fn, namespace=np) assert "_namespace = x + 1" in vec.source assert "xp = _namespace_1" in vec.source assert np.allclose(vec(np.asarray([1.0])), [4.0]) @@ -224,9 +230,9 @@ def test_user_local_named_namespace_is_kept() -> None: def test_fallback_with_pin() -> None: - fn = make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") + fn = support.make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") with pytest.warns(UserWarning, match="falling back") as record: - vec = vectorize(fn, fallback=True, namespace=np) + vec = array_vectorize.vectorize(fn, fallback=True, namespace=np) assert len(record) == 1 # warns once, at decoration time assert np.allclose(vec(np.asarray([5.0, -2.0])), [0.0, -2.0]) @@ -234,18 +240,18 @@ def test_fallback_with_pin() -> None: def test_fallback_with_namespace() -> None: import array_api_strict as xps - fn = make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") + fn = support.make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") with warnings.catch_warnings(): warnings.simplefilter("ignore") - vec = vectorize(fn, fallback=True) + vec = array_vectorize.vectorize(fn, fallback=True) pinned = vec.with_namespace(xps) assert list(map(float, pinned(xps.asarray([3.0])))) == [0.0] def test_fallback_pin_still_requires_array_arg() -> None: - fn = make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") + fn = support.make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(fn, fallback=True, namespace=np) + vec = array_vectorize.vectorize(fn, fallback=True, namespace=np) with pytest.raises(TypeError, match="requires at least one"): vec(3.0) @@ -254,13 +260,13 @@ def test_fallback_pin_still_requires_array_arg() -> None: def test_invalid_namespace_at_vectorize() -> None: - fn = make_fn(" return x + 1.0", defaults="x") + fn = support.make_fn(" return x + 1.0", defaults="x") with pytest.raises(TypeError, match="not an Array API namespace"): - vectorize(fn, namespace=object()) + array_vectorize.vectorize(fn, namespace=object()) def test_invalid_namespace_at_with_namespace() -> None: - vf = vectorize(make_fn(" return x + 1.0", defaults="x")) + vf = array_vectorize.vectorize(support.make_fn(" return x + 1.0", defaults="x")) with pytest.raises(TypeError, match="not an Array API namespace"): vf.with_namespace(object()) @@ -269,16 +275,18 @@ def test_invalid_namespace_at_with_namespace() -> None: def test_verify_with_pin_passes() -> None: - fn = make_fn(" if x < 0:\n return 0.0\n return x * math.exp(-x)", defaults="x") + fn = support.make_fn( + " if x < 0:\n return 0.0\n return x * math.exp(-x)", defaults="x" + ) xs = np.asarray([-1.0, 0.0, 2.0]) - vec = vectorize(fn, verify=(xs,), namespace=np) + vec = array_vectorize.vectorize(fn, verify=(xs,), namespace=np) assert np.allclose(vec(xs), np.where(xs < 0, 0.0, xs * np.exp(-xs))) def test_verify_with_pin_strict_backend() -> None: import array_api_strict as xps - fn = make_fn(" return x * 2.0", defaults="x") + fn = support.make_fn(" return x * 2.0", defaults="x") a = xps.asarray([1.0, 2.0]) - vec = vectorize(fn, verify=(a,), namespace=xps) + vec = array_vectorize.vectorize(fn, verify=(a,), namespace=xps) assert list(map(float, vec(a))) == [2.0, 4.0] From 7ad575ecd0c7f4c296941ae8d57f7b711c6e8ab7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 05:09:10 +0200 Subject: [PATCH 11/17] style: rewrite test and script imports to module-only form --- .claude/hooks/format_python.py | 4 +- examples/demo.py | 8 +- scripts/smoke.py | 6 +- tests/backends/test_array_api_strict.py | 122 +++--- tests/backends/test_jax_torch_compilation.py | 44 +- tests/behavior/test_arithmetic_dtypes.py | 320 ++++++++------- tests/behavior/test_calls_math_casts.py | 46 ++- tests/behavior/test_closures.py | 16 +- tests/behavior/test_control_flow.py | 26 +- tests/behavior/test_cse.py | 10 +- tests/behavior/test_divergences.py | 36 +- tests/behavior/test_docstrings.py | 18 +- tests/behavior/test_helpers.py | 48 +-- tests/behavior/test_inspectability.py | 34 +- tests/behavior/test_lambdas.py | 40 +- tests/behavior/test_loops.py | 178 ++++---- tests/behavior/test_memoization.py | 30 +- tests/behavior/test_name_allocation.py | 26 +- tests/behavior/test_public_api.py | 5 +- tests/conftest.py | 20 +- tests/corpus.py | 4 +- tests/differential/test_scalar_equivalence.py | 60 +-- tests/features/test_fallback.py | 22 +- tests/features/test_protect_domains.py | 78 ++-- tests/features/test_verify.py | 46 ++- tests/golden/test_golden.py | 17 +- tests/perf/test_bench.py | 15 +- tests/perf/test_speed_gate.py | 4 +- tests/support.py | 12 +- tests/test_agent_layout.py | 10 +- tests/unit/test_builtins.py | 21 +- tests/unit/test_codegen.py | 145 ++++--- tests/unit/test_docstring.py | 100 ++--- tests/unit/test_extract.py | 126 +++--- tests/unit/test_fuzz.py | 4 +- tests/unit/test_ir.py | 98 +++-- tests/unit/test_layering.py | 18 +- tests/unit/test_lower.py | 379 ++++++++++-------- tests/unit/test_optimize.py | 249 ++++++------ tests/unit/test_runtime.py | 64 +-- tests/unit/test_validate.py | 106 ++--- 41 files changed, 1378 insertions(+), 1237 deletions(-) diff --git a/.claude/hooks/format_python.py b/.claude/hooks/format_python.py index cc3f4bf..efa0e47 100755 --- a/.claude/hooks/format_python.py +++ b/.claude/hooks/format_python.py @@ -7,12 +7,12 @@ """ import json +import pathlib import subprocess import sys -from pathlib import Path path = json.load(sys.stdin).get("tool_input", {}).get("file_path", "") -if path.endswith((".py", ".pyi")) and Path(path).is_file(): +if path.endswith((".py", ".pyi")) and pathlib.Path(path).is_file(): ruff = ["uv", "run", "--frozen", "ruff"] subprocess.run([*ruff, "check", "--select", "I", "--fix", "--force-exclude", "-q", path]) subprocess.run([*ruff, "format", "--force-exclude", "-q", path]) diff --git a/examples/demo.py b/examples/demo.py index 2062f8b..7cd05a5 100644 --- a/examples/demo.py +++ b/examples/demo.py @@ -4,10 +4,10 @@ import numpy as np -from array_vectorize import vectorize +import array_vectorize -@vectorize +@array_vectorize.vectorize def sigmoid(x): """Numerically stable logistic sigmoid.""" if x >= 0: @@ -16,13 +16,13 @@ def sigmoid(x): return z / (1.0 + z) -@vectorize +@array_vectorize.vectorize def gelu(x): """GELU activation (tanh approximation).""" return 0.5 * x * (1.0 + math.tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))) -@vectorize +@array_vectorize.vectorize def l2_dist(x, y): """Distance accumulated in a constant-trip loop, then sqrt.""" d = 0.0 diff --git a/scripts/smoke.py b/scripts/smoke.py index 94ee023..f511f3a 100644 --- a/scripts/smoke.py +++ b/scripts/smoke.py @@ -4,10 +4,10 @@ import numpy as np -from array_vectorize import vectorize +import array_vectorize -@vectorize +@array_vectorize.vectorize def psi(x): """Exponential decay on the positive half, zero below it.""" if x < 0: @@ -15,7 +15,7 @@ def psi(x): return x * math.exp(-x) -@vectorize +@array_vectorize.vectorize def clamp(x, lo=0.0, hi=1.0): """Clamp x into [lo, hi].""" if x < lo: diff --git a/tests/backends/test_array_api_strict.py b/tests/backends/test_array_api_strict.py index 1814e57..d2f4520 100644 --- a/tests/backends/test_array_api_strict.py +++ b/tests/backends/test_array_api_strict.py @@ -7,9 +7,9 @@ import numpy as np import pytest -from support import make_fn, make_module, vfn +import support -from array_vectorize import vectorize +import array_vectorize # ----------------------------------------------- corpus and parameter preamble @@ -18,7 +18,7 @@ def test_backend_array_api_strict() -> None: import array_api_strict as xps a = xps.asarray([1.0, -2.0, 3.0]) - got = vfn("relu")(a) + got = support.vfn("relu")(a) assert isinstance(got, xps.asarray([1.0]).__class__) assert list(map(float, got)) == [1.0, 0.0, 3.0] @@ -28,7 +28,7 @@ def test_preamble_normalizes_scalar_params() -> None: # omitted defaults are raw Python scalars at runtime; the generated # preamble normalizes every parameter with xp.asarray - vec = vectorize(make_fn(" return x + y")) + vec = array_vectorize.vectorize(support.make_fn(" return x + y")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 3.0 @@ -36,7 +36,7 @@ def test_preamble_normalizes_scalar_params() -> None: def test_int_default_scalar_promotion_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return x + y", defaults="x, y=1")) + vec = array_vectorize.vectorize(support.make_fn(" return x + y", defaults="x, y=1")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 2.0 @@ -47,7 +47,7 @@ def test_int_default_scalar_promotion_strict() -> None: def test_constant_scalar_math_on_strict_backend() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return x + math.sqrt(4.0)")) + vec = array_vectorize.vectorize(support.make_fn(" return x + math.sqrt(4.0)")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 3.0 @@ -55,7 +55,7 @@ def test_constant_scalar_math_on_strict_backend() -> None: def test_literal_binding_math_call() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" y = 4\n return x + math.sqrt(y)")) + vec = array_vectorize.vectorize(support.make_fn(" y = 4\n return x + math.sqrt(y)")) assert np.allclose(vec(np.asarray([1.0])), [3.0]) assert float(vec(xps.asarray([1.0]))[0]) == 3.0 @@ -63,7 +63,7 @@ def test_literal_binding_math_call() -> None: def test_scalar_default_math_call_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return x + math.sqrt(y)")) + vec = array_vectorize.vectorize(support.make_fn(" return x + math.sqrt(y)")) assert np.allclose(vec(np.asarray([1.0])), [1.0 + math.sqrt(2.0)]) assert float(vec(xps.asarray([1.0]))[0]) == 1.0 + math.sqrt(2.0) @@ -71,7 +71,7 @@ def test_scalar_default_math_call_strict() -> None: def test_computed_scalar_math_arg_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return x + math.sqrt(y + 0.0)")) + vec = array_vectorize.vectorize(support.make_fn(" return x + math.sqrt(y + 0.0)")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 1.0 + math.sqrt(2.0) @@ -79,7 +79,9 @@ def test_computed_scalar_math_arg_strict() -> None: def test_computed_local_scalar_math_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = 1.0 + 3.0\n return x + math.sqrt(a)")) + vec = array_vectorize.vectorize( + support.make_fn(" a = 1.0 + 3.0\n return x + math.sqrt(a)") + ) got = vec(xps.asarray([1.0])) assert float(got[0]) == 3.0 @@ -87,8 +89,8 @@ def test_computed_local_scalar_math_strict() -> None: def test_loop_var_math_call_strict() -> None: import array_api_strict as xps - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " s = 0.0\n for i in range(3):\n s = s + math.sin(i)\n return s + x" ) ) @@ -99,8 +101,8 @@ def test_loop_var_math_call_strict() -> None: def test_loop_index_math_expression_strict() -> None: import array_api_strict as xps - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " s = 0.0\n for i in range(2):\n s = s + math.sqrt(i + 1.0)\n return s" ) ) @@ -112,8 +114,8 @@ def test_loop_index_math_expression_strict() -> None: def test_scalar_phi_strict() -> None: import array_api_strict as xps - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " b = 4.0\n" " for i in range(0):\n" " if x > 0:\n" @@ -132,18 +134,18 @@ def test_scalar_phi_strict() -> None: def test_verify_on_strict_backend() -> None: import array_api_strict as xps - mod = make_module("def subject(x):\n return x\n") - vec = vectorize(mod.subject, verify=(xps.asarray([1.0, 2.0]),)) + mod = support.make_module("def subject(x):\n return x\n") + vec = array_vectorize.vectorize(mod.subject, verify=(xps.asarray([1.0, 2.0]),)) assert list(map(float, vec(xps.asarray([1.0, 2.0])))) == [1.0, 2.0] def test_fallback_bool_arithmetic_on_strict_backend() -> None: import array_api_strict as xps - fn = make_fn(" while x > 0:\n x = x - 1\n b = x > 0\n return b + b") + fn = support.make_fn(" while x > 0:\n x = x - 1\n b = x > 0\n return b + b") with warnings.catch_warnings(): warnings.simplefilter("ignore") - vec = vectorize(fn, fallback=True) + vec = array_vectorize.vectorize(fn, fallback=True) got = vec(xps.asarray([3.0])) assert list(map(int, got)) == [0] @@ -154,7 +156,7 @@ def test_fallback_bool_arithmetic_on_strict_backend() -> None: def test_bool_arithmetic_strict_backend() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return (x > 0) + (x > 0)")) + vec = array_vectorize.vectorize(support.make_fn(" return (x > 0) + (x > 0)")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 2.0 @@ -162,7 +164,7 @@ def test_bool_arithmetic_strict_backend() -> None: def test_bool_mixed_arithmetic_strict_backend() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return (x > 0) + 1")) + vec = array_vectorize.vectorize(support.make_fn(" return (x > 0) + 1")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 2.0 @@ -170,7 +172,7 @@ def test_bool_mixed_arithmetic_strict_backend() -> None: def test_bool_plus_float_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return (x > 0) + x")) + vec = array_vectorize.vectorize(support.make_fn(" return (x > 0) + x")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 2.0 @@ -178,7 +180,7 @@ def test_bool_plus_float_strict() -> None: def test_bool_plus_float_literal_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return (x > 0) + 1.5")) + vec = array_vectorize.vectorize(support.make_fn(" return (x > 0) + 1.5")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 2.5 @@ -186,7 +188,7 @@ def test_bool_plus_float_literal_strict() -> None: def test_arith_result_negative_power_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = (x > 0) + 1\n return a ** -1")) + vec = array_vectorize.vectorize(support.make_fn(" a = (x > 0) + 1\n return a ** -1")) assert float(vec(xps.asarray([2]))[0]) == 0.5 @@ -196,7 +198,7 @@ def test_arith_result_negative_power_strict() -> None: def test_min_max_literal_args_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return min(x, 3) + max(y, 1)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 3) + max(y, 1)")) got = vec(xps.asarray([1.0])) assert float(got[0]) == 3.0 # min(1, 3) + max(2, 1) == 1 + 2 @@ -204,7 +206,7 @@ def test_min_max_literal_args_strict() -> None: def test_min_literal_strict_int_input() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return min(x, 3)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 3)")) got = vec(xps.asarray([1], dtype=xps.int64)) assert int(got[0]) == 1 @@ -212,7 +214,7 @@ def test_min_literal_strict_int_input() -> None: def test_min_literal_strict_float_default_param() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return min(y, 3)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(y, 3)")) got = vec(xps.asarray([1.0])) # constant w.r.t. x: 0-d result is broadcast for comparison assert np.allclose(np.broadcast_to(np.asarray(got, dtype=float), (1,)), [2.0]) @@ -221,7 +223,9 @@ def test_min_literal_strict_float_default_param() -> None: def test_all_literal_minmax_refs_fold() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = 1.0\n b = 2.0\n return x + min(a, b)")) + vec = array_vectorize.vectorize( + support.make_fn(" a = 1.0\n b = 2.0\n return x + min(a, b)") + ) got = vec(xps.asarray([1.0])) assert float(got[0]) == 2.0 @@ -229,14 +233,14 @@ def test_all_literal_minmax_refs_fold() -> None: def test_minmax_promotion_matrix_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return min(x, 3)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 3)")) assert int(vec(xps.asarray([1], dtype=xps.int64))[0]) == 1 - vec = vectorize(make_fn(" return min(x, 1.5) + max(y, 1)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 1.5) + max(y, 1)")) got = vec(xps.asarray([2.0])) assert float(got[0]) == 3.5 # min(2, 1.5) + max(2, 1) - vec = vectorize(make_fn(" return min(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, y)")) got = vec(xps.asarray([1.0, 3.0]), xps.asarray([2.0, 2.0])) assert np.allclose(np.asarray(got, dtype=float), [1.0, 2.0]) @@ -246,7 +250,7 @@ def test_max_variadic_literals_strict() -> None: # 1 fits uint8 but 1.5 forces float64: BOTH literals must cast to # float64 (a uint8 literal would break strict same-dtype promotion) - vec = vectorize(make_fn(" return max(x, 1, 1.5)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 1, 1.5)")) got = vec(xps.asarray([2], dtype=xps.uint8)) assert int(got[0]) == 2 @@ -254,7 +258,7 @@ def test_max_variadic_literals_strict() -> None: def test_bool_minmax_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return min(x > 0, x > 1)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x > 0, x > 1)")) got = vec(xps.asarray([2, 0])) assert list(np.asarray(got, dtype=bool)) == [True, False] @@ -262,7 +266,7 @@ def test_bool_minmax_strict() -> None: def test_bool_minmax_strict_max() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return max(x > 0, x > 1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x > 0, x > 1)")) got = vec(xps.asarray([2, 0])) assert list(np.asarray(got, dtype=bool)) == [True, False] @@ -270,7 +274,7 @@ def test_bool_minmax_strict_max() -> None: def test_bool_only_minmax_bitwise_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return min(x > 0, x > 1) & (x > 0)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x > 0, x > 1) & (x > 0)")) got = vec(xps.asarray([2])) assert bool(np.asarray(got).reshape(-1)[0]) @@ -281,7 +285,7 @@ def test_bool_only_minmax_bitwise_strict() -> None: def test_minmax_runtime_bool_arithmetic_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = min(x, True)\n return a + a")) + vec = array_vectorize.vectorize(support.make_fn(" a = min(x, True)\n return a + a")) got = vec(xps.asarray([True, False])) assert list(map(int, np.asarray(got))) == [2, 0] @@ -289,7 +293,7 @@ def test_minmax_runtime_bool_arithmetic_strict() -> None: def test_unary_negate_maybe_bool_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = min(x, True)\n return -a")) + vec = array_vectorize.vectorize(support.make_fn(" a = min(x, True)\n return -a")) got = vec(xps.asarray([True, False])) assert list(map(int, np.asarray(got))) == [-1, 0] @@ -297,7 +301,9 @@ def test_unary_negate_maybe_bool_strict() -> None: def test_bitwise_maybe_bool_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = min(x, True)\n b = a & True\n return b + b")) + vec = array_vectorize.vectorize( + support.make_fn(" a = min(x, True)\n b = a & True\n return b + b") + ) got = vec(xps.asarray([True, False])) assert list(map(int, np.asarray(got))) == [2, 0] @@ -305,7 +311,7 @@ def test_bitwise_maybe_bool_strict() -> None: def test_truediv_bool_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return min(x, True) / 2")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, True) / 2")) got = vec(xps.asarray([True, False])) assert list(map(float, np.asarray(got))) == [0.5, 0.0] @@ -316,7 +322,9 @@ def test_truediv_bool_strict() -> None: def test_uint64_plus_nonnegative_int64_array_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = max(x, True)\n return a + y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a + y", defaults="x, y") + ) got = vec(xps.asarray([2**63 + 1], dtype=xps.uint64), xps.asarray([0], dtype=xps.int64)) assert int(got[0]) == 9223372036854775809 @@ -324,7 +332,9 @@ def test_uint64_plus_nonnegative_int64_array_strict() -> None: def test_uint64_minus_uint64_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = max(x, True)\n return a - y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a - y", defaults="x, y") + ) got = vec(xps.asarray([5], dtype=xps.uint64), xps.asarray([7], dtype=xps.uint64)) assert int(got[0]) == -2 @@ -332,7 +342,9 @@ def test_uint64_minus_uint64_strict() -> None: def test_int64_minus_uint64_huge_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = max(x, True)\n return y - a", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return y - a", defaults="x, y") + ) got = vec(xps.asarray([2**63], dtype=xps.uint64), xps.asarray([2**63 - 1], dtype=xps.int64)) assert int(got[0]) == -1 @@ -340,7 +352,7 @@ def test_int64_minus_uint64_huge_strict() -> None: def test_literal_minus_uint64_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return 0 - max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return 0 - max(x, True)")) got = vec(xps.asarray([1], dtype=xps.uint64)) assert int(got[0]) == -1 @@ -348,7 +360,7 @@ def test_literal_minus_uint64_strict() -> None: def test_uint64_minus_negative_literal_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return max(x, True) - -2")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, True) - -2")) got = vec(xps.asarray([2**63 + 3], dtype=xps.uint64)) assert int(got[0]) == 2**63 + 5 @@ -356,7 +368,7 @@ def test_uint64_minus_negative_literal_strict() -> None: def test_negative_literal_minus_uint64_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return -2 - max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return -2 - max(x, True)")) got = vec(xps.asarray([5], dtype=xps.uint64)) assert int(got[0]) == -7 @@ -367,7 +379,7 @@ def test_negative_literal_minus_uint64_strict() -> None: def test_uint64_floordiv_negative_divisor_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return max(x, True) // -3")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, True) // -3")) got = vec(xps.asarray([5], dtype=xps.uint64)) assert int(got[0]) == -2 @@ -375,7 +387,7 @@ def test_uint64_floordiv_negative_divisor_strict() -> None: def test_uint64_negative_literal_floordiv_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" return -3 // max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return -3 // max(x, True)")) got = vec(xps.asarray([5], dtype=xps.uint64)) assert int(got[0]) == -1 @@ -383,7 +395,9 @@ def test_uint64_negative_literal_floordiv_strict() -> None: def test_uint64_floordiv_positive_int64_array_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = max(x, True)\n return a // y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a // y", defaults="x, y") + ) got = vec(xps.asarray([2**63 + 3], dtype=xps.uint64), xps.asarray([2], dtype=xps.int64)) assert int(got[0]) == 4611686018427387905 @@ -391,7 +405,9 @@ def test_uint64_floordiv_positive_int64_array_strict() -> None: def test_uint64_mod_positive_int64_array_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = max(x, True)\n return a % y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a % y", defaults="x, y") + ) got = vec(xps.asarray([2**63 + 3], dtype=xps.uint64), xps.asarray([2], dtype=xps.int64)) assert int(got[0]) == 1 @@ -399,13 +415,13 @@ def test_uint64_mod_positive_int64_array_strict() -> None: def test_uint64_mod_int64_min_divisor_strict() -> None: import array_api_strict as xps - mod = make_module( + mod = support.make_module( "import math\n\n" "BOUND = -(2**63)\n\n" "def subject(x, y=2.0):\n" " return max(x, True) % BOUND\n" ) - vec = vectorize(mod.subject) + vec = array_vectorize.vectorize(mod.subject) got = vec(xps.asarray([1], dtype=xps.uint64)) assert int(got[0]) == -9223372036854775807 @@ -413,6 +429,6 @@ def test_uint64_mod_int64_min_divisor_strict() -> None: def test_uint64_truediv_negative_divisor_strict() -> None: import array_api_strict as xps - vec = vectorize(make_fn(" a = max(x, True)\n return a / -2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a / -2")) got = vec(xps.asarray([3], dtype=xps.uint64)) assert float(got[0]) == -1.5 diff --git a/tests/backends/test_jax_torch_compilation.py b/tests/backends/test_jax_torch_compilation.py index 9d1c981..51743e6 100644 --- a/tests/backends/test_jax_torch_compilation.py +++ b/tests/backends/test_jax_torch_compilation.py @@ -22,18 +22,18 @@ from __future__ import annotations import importlib.util +import pathlib from collections.abc import Callable, Sequence -from pathlib import Path from typing import Any import numpy as np import pytest -from support import make_fn, make_module +import support -from array_vectorize import vectorize +import array_vectorize spec = importlib.util.spec_from_file_location( - "vec_corpus_c", Path(__file__).parent.parent / "corpus.py" + "vec_corpus_c", pathlib.Path(__file__).parent.parent / "corpus.py" ) assert spec is not None assert spec.loader is not None @@ -107,8 +107,8 @@ def test_jax_jit_detected_matches_eager(name: str) -> None: fn = getattr(CORPUS, name) args, kwargs = _call_args(name, jnp.asarray) - eager = vectorize(fn)(*args, **kwargs) - got = jax.jit(vectorize(fn))(*args, **kwargs) + eager = array_vectorize.vectorize(fn)(*args, **kwargs) + got = jax.jit(array_vectorize.vectorize(fn))(*args, **kwargs) assert np.allclose(np.asarray(got), np.asarray(eager), equal_nan=True) @@ -120,8 +120,8 @@ def test_jax_jit_pinned_matches_eager(name: str) -> None: fn = getattr(CORPUS, name) args, kwargs = _call_args(name, jnp.asarray) - eager = vectorize(fn)(*args, **kwargs) - got = jax.jit(vectorize(fn, namespace=jnp))(*args, **kwargs) + eager = array_vectorize.vectorize(fn)(*args, **kwargs) + got = jax.jit(array_vectorize.vectorize(fn, namespace=jnp))(*args, **kwargs) assert np.allclose(np.asarray(got), np.asarray(eager), equal_nan=True) @@ -144,8 +144,8 @@ def test_torch_compile_detected_matches_eager(name: str) -> None: fn = getattr(CORPUS, name) args, kwargs = _call_args(name, _torch_conv) - eager = vectorize(fn)(*args, **kwargs) - got = torch.compile(vectorize(fn))(*args, **kwargs) + eager = array_vectorize.vectorize(fn)(*args, **kwargs) + got = torch.compile(array_vectorize.vectorize(fn))(*args, **kwargs) assert torch.allclose(got, eager, equal_nan=True) @@ -158,8 +158,10 @@ def test_torch_compile_pinned_matches_eager(name: str) -> None: fn = getattr(CORPUS, name) args, kwargs = _call_args(name, _torch_conv) - eager = vectorize(fn, namespace=array_api_compat.torch)(*args, **kwargs) - got = torch.compile(vectorize(fn, namespace=array_api_compat.torch))(*args, **kwargs) + eager = array_vectorize.vectorize(fn, namespace=array_api_compat.torch)(*args, **kwargs) + got = torch.compile(array_vectorize.vectorize(fn, namespace=array_api_compat.torch))( + *args, **kwargs + ) assert torch.allclose(got, eager, equal_nan=True) @@ -178,7 +180,7 @@ def test_verify_with_torch_args() -> None: t = torch.tensor(X, dtype=torch.float64) # previously: VectorizationError("verify= needs at least one Array API # array in the example inputs") - vec = vectorize(CORPUS.psi, verify=(t,)) + vec = array_vectorize.vectorize(CORPUS.psi, verify=(t,)) got = vec(t) expected = torch.where(t < 0, torch.tensor(0.0), t * torch.exp(-t)) assert torch.allclose(got, expected) @@ -188,9 +190,9 @@ def test_fallback_with_torch_tensors() -> None: pytest.importorskip("torch") import torch - fn = make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") + fn = support.make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(fn, fallback=True) + vec = array_vectorize.vectorize(fn, fallback=True) # previously: TypeError("vectorized fallback ... requires at least one # Array API array argument") t = torch.tensor([5.0, -2.0], dtype=torch.float64) @@ -207,9 +209,9 @@ def test_fallback_pinned_with_torch_tensors() -> None: import array_api_compat.torch import torch - fn = make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") + fn = support.make_fn(" while x > 0:\n x = x - 1\n return x", defaults="x") with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(fn, fallback=True, namespace=array_api_compat.torch) + vec = array_vectorize.vectorize(fn, fallback=True, namespace=array_api_compat.torch) t = torch.tensor([5.0, -2.0], dtype=torch.float64) got = vec(t) assert isinstance(got, torch.Tensor) @@ -220,7 +222,7 @@ def test_closure_captured_torch_tensor() -> None: pytest.importorskip("torch") import torch - mod = make_module( + mod = support.make_module( "import torch\n" "T = torch.asarray([1.0, 2.0, 3.0], dtype=torch.float64)\n" "\n" @@ -230,7 +232,7 @@ def test_closure_captured_torch_tensor() -> None: ) # previously: VectorizationError (closure/global 'T' has unsupported # type 'Tensor'); now a hidden kw-only closure param like numpy captures - vec = vectorize(mod.subject) + vec = array_vectorize.vectorize(mod.subject) assert "def subject_vec(x, *, T=None):" in vec.source assert vec.__kwdefaults__["T"] is mod.T t = torch.tensor([10.0, 20.0, 30.0], dtype=torch.float64) @@ -241,7 +243,7 @@ def test_closure_captured_torch_tensor() -> None: def test_closure_captured_jax_array() -> None: jax = pytest.importorskip("jax") jax.config.update("jax_enable_x64", True) - mod = make_module( + mod = support.make_module( "import jax.numpy as jnp\n" "T = jnp.asarray([1.0, 2.0, 3.0])\n" "\n" @@ -249,7 +251,7 @@ def test_closure_captured_jax_array() -> None: "def subject(x):\n" " return x + T\n" ) - vec = vectorize(mod.subject) + vec = array_vectorize.vectorize(mod.subject) assert "def subject_vec(x, *, T=None):" in vec.source assert vec.__kwdefaults__["T"] is mod.T got = vec(mod.T * 10.0) diff --git a/tests/behavior/test_arithmetic_dtypes.py b/tests/behavior/test_arithmetic_dtypes.py index 60389a1..d90376c 100644 --- a/tests/behavior/test_arithmetic_dtypes.py +++ b/tests/behavior/test_arithmetic_dtypes.py @@ -6,9 +6,9 @@ import numpy as np import pytest -from support import make_fn, make_module, vec_of, vfn +import support -from array_vectorize import vectorize +import array_vectorize X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) Y = np.asarray([3.0, 1.5, 1.0, 0.5, -1.0, -3.0]) @@ -22,60 +22,60 @@ def py_scalar(got: Any) -> int: def test_add() -> None: - assert np.allclose(vfn("add")(X, Y), X + Y) + assert np.allclose(support.vfn("add")(X, Y), X + Y) def test_arith_ops() -> None: expected = (X + Y - X * Y) / (X * X + Y * Y) - assert np.allclose(vfn("arith_ops")(X, Y), expected) + assert np.allclose(support.vfn("arith_ops")(X, Y), expected) def test_bitwise_ops() -> None: xi = np.asarray([-3, 2, 5, 7], dtype=np.int64) yi = np.asarray([4, 1, 6, 2], dtype=np.int64) expected = (xi & yi) | (xi ^ 3) << 1 >> 2 - assert np.array_equal(vfn("bitwise_ops")(xi, yi), expected) + assert np.array_equal(support.vfn("bitwise_ops")(xi, yi), expected) def test_unary_ops_int() -> None: xi = np.asarray([-3, 2, 5], dtype=np.int64) - assert np.array_equal(vfn("unary_ops")(xi), -xi + (+xi) - (~xi)) + assert np.array_equal(support.vfn("unary_ops")(xi), -xi + (+xi) - (~xi)) def test_minmax() -> None: expected = np.minimum(X, Y) + np.maximum(np.maximum(X, Y), 3.0) - assert np.allclose(vfn("minmax")(X, Y), expected) + assert np.allclose(support.vfn("minmax")(X, Y), expected) def test_compare_chain() -> None: - got = vfn("compare_chain")(X) + got = support.vfn("compare_chain")(X) assert got.dtype == np.bool_ assert np.array_equal(got, (X > 0) & (X < 1)) def test_boolop_logical() -> None: - got = vfn("boolop_logical")(X, Y) + got = support.vfn("boolop_logical")(X, Y) expected = ((X > 0) & (Y > 0)) | ~(X < 0) assert np.array_equal(got, expected) def test_boolop_numeric_exact_including_nan() -> None: arr = np.asarray([-0.0, 0.0, 2.0, np.nan]) - got = vfn("boolop_numeric")(arr, np.asarray([5.0, 5.0, 0.0, 1.0])) + got = support.vfn("boolop_numeric")(arr, np.asarray([5.0, 5.0, 0.0, 1.0])) # Python: x and y -> x if falsy else y; NaN is truthy expected = np.asarray([arr[0], 0.0, 0.0, 1.0]) assert np.array_equal(got, expected, equal_nan=True) def test_clamp() -> None: - assert np.allclose(vfn("clamp")(X), np.clip(X, 0.0, 1.0)) - assert np.allclose(vfn("clamp")(X, -1.0, 0.5), np.clip(X, -1.0, 0.5)) - assert np.allclose(vfn("clamp")(X, hi=0.25), np.clip(X, 0.0, 0.25)) + assert np.allclose(support.vfn("clamp")(X), np.clip(X, 0.0, 1.0)) + assert np.allclose(support.vfn("clamp")(X, -1.0, 0.5), np.clip(X, -1.0, 0.5)) + assert np.allclose(support.vfn("clamp")(X, hi=0.25), np.clip(X, 0.0, 0.25)) def test_scalar_python_args_accepted() -> None: # Python scalars ride along via standard promotion - assert np.allclose(vfn("add")(X, 2.0), X + 2.0) + assert np.allclose(support.vfn("add")(X, 2.0), X + 2.0) # ------------------------------------------------------------- bool arithmetic @@ -90,48 +90,50 @@ def test_scalar_python_args_accepted() -> None: ], ) def test_boolean_arithmetic_through_assignments(body: str) -> None: - vec, fn = vec_of(body) + vec, fn = support.vec_of(body) xs = np.asarray([-1.0, 1.0]) assert list(map(int, vec(xs))) == [int(fn(float(v))) for v in xs] def test_unary_negate_bool() -> None: # -True is -1 in Python; backends cannot negate bool arrays - vec = vectorize(make_fn(" return -(x > 0)")) + vec = array_vectorize.vectorize(support.make_fn(" return -(x > 0)")) assert vec(np.asarray([2.0]))[0] == -1 def test_bool_plus_large_int_exact() -> None: - vec = vectorize(make_fn(" return (x > 0) + 9007199254740992")) + vec = array_vectorize.vectorize(support.make_fn(" return (x > 0) + 9007199254740992")) got = vec(np.asarray([1.0])) assert got[0] == 9007199254740993 def test_bool_arith_then_bitwise() -> None: - vec = vectorize(make_fn(" return ((x > 0) + (x > 0)) & 1")) + vec = array_vectorize.vectorize(support.make_fn(" return ((x > 0) + (x > 0)) & 1")) got = vec(np.asarray([1.0])) assert got[0] == 0 def test_bool_arith_times_hundred() -> None: - vec = vectorize(make_fn(" return ((x > 0) + (x > 0)) * 100")) + vec = array_vectorize.vectorize(support.make_fn(" return ((x > 0) + (x > 0)) * 100")) assert vec(np.asarray([1.0]))[0] == 200 def test_bool_plus_int8_array_headroom() -> None: - vec = vectorize(make_fn(" return (x > 0) + y", defaults="x, y")) + vec = array_vectorize.vectorize(support.make_fn(" return (x > 0) + y", defaults="x, y")) got = vec(np.asarray([1.0]), np.asarray([100], dtype=np.int8)) assert got[0] == 101 def test_bool_bitwise_stays_bool() -> None: - vec = vectorize(make_fn(" b = (x > 0) & (x < 3)\n return b + b")) + vec = array_vectorize.vectorize(support.make_fn(" b = (x > 0) & (x < 3)\n return b + b")) got = vec(np.asarray([1.0])) assert got[0] == 2 def test_bool_bitwise_or_xor_stay_bool() -> None: - vec = vectorize(make_fn(" return ((x > 0) | (x > 1)) + ((x > 0) ^ (x > 1))")) + vec = array_vectorize.vectorize( + support.make_fn(" return ((x > 0) | (x > 1)) + ((x > 0) ^ (x > 1))") + ) got = vec(np.asarray([1.0])) assert got[0] == 2 @@ -141,12 +143,12 @@ def test_bool_bitwise_or_xor_stay_bool() -> None: def test_pow_negative_exponent_int_base() -> None: # Python promotes int ** negative-int to float; arrays raise instead - vec = vectorize(make_fn(" return x ** -1")) + vec = array_vectorize.vectorize(support.make_fn(" return x ** -1")) assert np.allclose(vec(np.asarray([2.0, 4.0])), [0.5, 0.25]) def test_pow_negative_exponent_literal_base() -> None: - vec = vectorize(make_fn(" return 2 ** -1")) + vec = array_vectorize.vectorize(support.make_fn(" return 2 ** -1")) got = vec(np.asarray([0.0])) # constant functions fold to a scalar result; broadcast for comparison assert np.allclose(np.broadcast_to(np.asarray(got), (1,)), [0.5]) @@ -155,13 +157,13 @@ def test_pow_negative_exponent_literal_base() -> None: def test_arith_result_negative_power() -> None: # (x > 0) + 1 is provably int: the negative-exponent float cast # must fire on the helper result - vec = vectorize(make_fn(" a = (x > 0) + 1\n return a ** -1")) + vec = array_vectorize.vectorize(support.make_fn(" a = (x > 0) + 1\n return a ** -1")) assert vec(np.asarray([2]))[0] == 0.5 def test_arith_result_float_kind() -> None: # float operands keep float: no int-power issue to begin with - vec = vectorize(make_fn(" a = (x > 0) + 0.5\n return a ** -1")) + vec = array_vectorize.vectorize(support.make_fn(" a = (x > 0) + 0.5\n return a ** -1")) assert vec(np.asarray([2]))[0] == 1 / 1.5 # True + 0.5 == 1.5 @@ -169,56 +171,58 @@ def test_arith_result_float_kind() -> None: def test_minmax_all_literals_fold() -> None: - vec = vectorize(make_fn(" return min(3, 5) + x")) + vec = array_vectorize.vectorize(support.make_fn(" return min(3, 5) + x")) assert np.allclose(vec(np.asarray([1.0])), [4.0]) def test_minmax_loop_var_sibling() -> None: # the loop variable is a raw Python int per iteration: min/max must # wrap it (asarray) and give int literals a matching dtype - vec = vectorize( - make_fn(" m = x\n for i in range(4):\n m = min(i, 2) + m\n return m") + vec = array_vectorize.vectorize( + support.make_fn( + " m = x\n for i in range(4):\n m = min(i, 2) + m\n return m" + ) ) # min(i, 2) over i = 0..3 sums to 0 + 1 + 2 + 2 = 5 assert vec(np.asarray([0.0]))[0] == 5.0 def test_negative_literal_other_contexts() -> None: - vec = vectorize(make_fn(" return min(x, -1) + max(y, -2.5)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, -1) + max(y, -2.5)")) got = vec(np.asarray([3.0]), np.asarray([-3.0])) assert got[0] == -3.5 # min(3, -1) + max(-3, -2.5) = -1 + -2.5 def test_min_literal_exact_large_int() -> None: - vec = vectorize(make_fn(" return min(x, 9007199254740993)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 9007199254740993)")) got = vec(np.asarray([9007199254740994], dtype=np.int64)) assert got[0] == 9007199254740993 def test_min_compound_sibling_exact_big_int() -> None: - vec = vectorize(make_fn(" return min(x + 0, 9007199254740993)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x + 0, 9007199254740993)")) got = vec(np.asarray([9007199254740994], dtype=np.int64)) assert got[0] == 9007199254740993 def test_min_float_literal_int_input() -> None: # the float bound must widen the common dtype, not truncate to int64 - vec = vectorize(make_fn(" return min(x, 1.5)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 1.5)")) got = vec(np.asarray([2], dtype=np.int64)) assert got[0] == 1.5 def test_max_literal_int8_no_overflow() -> None: # 300 must widen int8, not wrap to 44 - vec = vectorize(make_fn(" return max(x, 300)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 300)")) got = vec(np.asarray([1], dtype=np.int8)) assert got[0] == 300 def test_min_int_literal_fit_boundaries() -> None: - vec = vectorize(make_fn(" return min(x, 127)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 127)")) assert vec(np.asarray([100], dtype=np.int8))[0] == 100 - vec = vectorize(make_fn(" return max(x, 128)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 128)")) # 128 does not fit int8: promotion widens instead of wrapping assert vec(np.asarray([1], dtype=np.int8))[0] == 128 @@ -226,7 +230,7 @@ def test_min_int_literal_fit_boundaries() -> None: def test_max_bool_array_vs_literal() -> None: # bool and the literal 2 share the promotion class (int, 8): the fast # path must not pick bool and collapse 2 to True - vec = vectorize(make_fn(" return max(x, 2)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 2)")) got = vec(np.asarray([False])) assert got[0] == 2 @@ -234,13 +238,13 @@ def test_max_bool_array_vs_literal() -> None: def test_min_float32_literal_needs_wider_dtype() -> None: # 16777217 is not exactly representable in float32: the fast path # must reject it and promote to float64 - vec = vectorize(make_fn(" return min(x, 16777217)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 16777217)")) got = vec(np.asarray([3e7], dtype=np.float32)) assert got[0] == 16777217 def test_min_float32_exact_literal_keeps_float32() -> None: - vec = vectorize(make_fn(" return min(x, 1.5)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 1.5)")) got = vec(np.asarray([2.0], dtype=np.float32)) assert got[0] == 1.5 @@ -248,52 +252,52 @@ def test_min_float32_exact_literal_keeps_float32() -> None: def test_max_float16_literal_precision() -> None: # 2049 is not representable in float16: the fast path must reject it # and promote to float64 instead of rounding the bound to 2048 - vec = vectorize(make_fn(" return max(x, 2049.0)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 2049.0)")) got = vec(np.asarray([0.0], dtype=np.float16)) assert got[0] == 2049.0 def test_max_float16_exact_literal_keeps_dtype() -> None: - vec = vectorize(make_fn(" return max(x, 2048.0)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 2048.0)")) got = vec(np.asarray([0.0], dtype=np.float16)) assert got[0] == 2048.0 def test_min_float16_overflow_literal_promotes() -> None: # 1e10 overflows float16 in the IEEE round trip: promote to float64 - vec = vectorize(make_fn(" return min(x, 1e10)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 1e10)")) got = vec(np.asarray([0.0], dtype=np.float16)) assert got[0] == 0.0 def test_min_uint8_too_large_bound_clamps() -> None: # 300 never wins a min over uint8 values (max 255): clamp to 255 - vec = vectorize(make_fn(" return min(x, 300)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 300)")) got = vec(np.asarray([1], dtype=np.uint8)) assert got[0] == 1 def test_max_uint8_negative_bound_clamps() -> None: - vec = vectorize(make_fn(" return max(x, -1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, -1)")) got = vec(np.asarray([7], dtype=np.uint8)) assert got[0] == 7 def test_min_variadic_negative_and_float_literals() -> None: - vec = vectorize(make_fn(" return min(x, -1, 1.5)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, -1, 1.5)")) got = vec(np.asarray([2], dtype=np.uint8)) assert int(got[0]) == -1 def test_max_variadic_fitting_literals_keep_dtype() -> None: - vec = vectorize(make_fn(" return max(x, 1, 2)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 1, 2)")) got = vec(np.asarray([0], dtype=np.uint8)) assert int(got[0]) == 2 def test_fit_literal_alongside_clamping_literal() -> None: # 5 fits uint8 while -1 clamps: both stay in the shared dtype - vec = vectorize(make_fn(" return max(x, 5, -1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 5, -1)")) got = vec(np.asarray([1], dtype=np.uint8)) assert py_scalar(got) == 5 @@ -302,13 +306,13 @@ def test_fit_literal_alongside_clamping_literal() -> None: def test_minmax_identical_dtypes_passthrough() -> None: - vec = vectorize(make_fn(" return min(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, y)")) got = vec(np.asarray([3, 1], dtype=np.int8), np.asarray([2, 2], dtype=np.int8)) assert list(np.asarray(got)) == [2, 1] def test_min_float32_vs_int64_exact() -> None: - vec = vectorize(make_fn(" return min(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, y)")) got = vec(np.asarray([20000000], dtype=np.float32), np.asarray([16777217], dtype=np.int64)) assert got[0] == 16777217 @@ -316,13 +320,13 @@ def test_min_float32_vs_int64_exact() -> None: def test_max_uint8_int16_no_narrowing() -> None: # uint8 and int16 share the promotion class (int, 16): the fast path # must not pick uint8 and wrap 300 to 44 - vec = vectorize(make_fn(" return max(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, y)")) got = vec(np.asarray([1], dtype=np.uint8), np.asarray([300], dtype=np.int16)) assert got[0] == 300 def test_min_uint8_int16_negative() -> None: - vec = vectorize(make_fn(" return min(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, y)")) got = vec(np.asarray([1], dtype=np.uint8), np.asarray([-1], dtype=np.int16)) assert got[0] == -1 @@ -330,7 +334,7 @@ def test_min_uint8_int16_negative() -> None: def test_min_uint64_vs_float64() -> None: # uint64 and float64 share the promotion class (float, 64): the fast # path must not pick uint64 and truncate 1.5 to 1 - vec = vectorize(make_fn(" return min(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, y)")) got = vec(np.asarray([2**63], dtype=np.uint64), np.asarray([1.5])) assert got[0] == 1.5 @@ -338,17 +342,17 @@ def test_min_uint64_vs_float64() -> None: def test_minmax_promotion_matrix_exact() -> None: # one pass over the min/max promotion matrix (uint64 bounds, uint8 vs # int16, float literal vs int input, float16 precision): all exact - vec = vectorize(make_fn(" return max(x, -1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, -1)")) assert py_scalar(vec(np.asarray([2**63 + 1], dtype=np.uint64))) == 9223372036854775809 - vec = vectorize(make_fn(" return max(x, 1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 1)")) assert py_scalar(vec(np.asarray([9223372036854775809], dtype=np.uint64))) == 9223372036854775809 - vec = vectorize(make_fn(" return min(x, 300)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 300)")) assert py_scalar(vec(np.asarray([1], dtype=np.uint8))) == 1 - vec = vectorize(make_fn(" return max(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, y)")) assert py_scalar(vec(np.asarray([1], dtype=np.uint8), np.asarray([300], dtype=np.int16))) == 300 - vec = vectorize(make_fn(" return min(x, 1.5)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 1.5)")) assert vec(np.asarray([2], dtype=np.int64))[0] == 1.5 - vec = vectorize(make_fn(" return max(x, 2049.0)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 2049.0)")) assert vec(np.asarray([0.0], dtype=np.float16))[0] == 2049.0 @@ -356,7 +360,7 @@ def test_minmax_promotion_matrix_exact() -> None: def test_min_uint64_vs_literal() -> None: - vec = vectorize(make_fn(" return min(x, 1)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, 1)")) got = vec(np.asarray([2**63], dtype=np.uint64)) assert got[0] == 1 @@ -364,7 +368,7 @@ def test_min_uint64_vs_literal() -> None: def test_max_uint64_literal_exact() -> None: # the literal 1 fits uint64, so the fast path keeps uint64: the # result is exact, not the float64 rounding of 9223372036854775809 - vec = vectorize(make_fn(" return max(x, 1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, 1)")) got = vec(np.asarray([9223372036854775809], dtype=np.uint64)) assert got[0] == 9223372036854775809 @@ -374,7 +378,7 @@ def test_max_uint64_negative_literal_exact() -> None: # clamped into the unsigned dtype instead of forcing float64 (which # would round 9223372036854775809 to ...808). `-1` lowers to # UnaryOp(neg, 1) and must still be recognized as a literal. - vec = vectorize(make_fn(" return max(x, -1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, -1)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) # exact comparison via Python int (numpy scalar == rounds floats) assert int(got[0]) == 9223372036854775809 @@ -382,7 +386,7 @@ def test_max_uint64_negative_literal_exact() -> None: def test_min_uint64_negative_literal_wins_exactly() -> None: # min selects the negative bound: it must survive exactly - vec = vectorize(make_fn(" return min(x, -1)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, -1)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) assert int(got[0]) == -1 @@ -390,61 +394,61 @@ def test_min_uint64_negative_literal_wins_exactly() -> None: def test_min_uint64_always_winning_literal_exact() -> None: # the negative bound always wins: the result is the literal, exactly # representable in int64 (float64 would round -9007199254740993) - vec = vectorize(make_fn(" return min(x, -9007199254740993)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, -9007199254740993)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) assert py_scalar(got) == -9007199254740993 def test_min_uint64_result_supports_bitwise() -> None: # the min result must stay INTEGER so downstream int ops work - vec = vectorize(make_fn(" return min(x, -1) & 1")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, -1) & 1")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) assert py_scalar(got) == 1 def test_min_uint64_signed_positive_literals() -> None: # min results fit int64 when something signed can win - vec = vectorize(make_fn(" return min(x, -1, 5)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, -1, 5)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) assert py_scalar(got) == -1 def test_max_uint64_int64_arrays_exact() -> None: - vec = vectorize(make_fn(" return max(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, y)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([-1], dtype=np.int64)) assert py_scalar(got) == 2**63 + 1 def test_min_uint64_int64_arrays_exact() -> None: - vec = vectorize(make_fn(" return min(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, y)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([-1], dtype=np.int64)) assert py_scalar(got) == -1 def test_max_uint64_all_unsigned_widens() -> None: - vec = vectorize(make_fn(" return max(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, y)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([7], dtype=np.uint8)) assert py_scalar(got) == 2**63 + 1 def test_max_uint64_mixed_arrays_negative_literal() -> None: - vec = vectorize(make_fn(" return max(x, y, -1)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, y, -1)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([-5], dtype=np.int64)) assert py_scalar(got) == 2**63 + 1 def test_max_uint64_mixed_arrays_positive_literal() -> None: - vec = vectorize(make_fn(" return max(x, y, 3)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, y, 3)")) got = vec(np.asarray([1], dtype=np.uint64), np.asarray([-5], dtype=np.int64)) assert py_scalar(got) == 3 def test_uint64_bool_mix_minmax() -> None: # the uint64 paths' signed-clamp comparisons must never see bools - vec = vectorize(make_fn(" return min(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, y)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([True])) assert int(got[0]) == 1 - vec = vectorize(make_fn(" return max(x, y)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, y)")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([True])) assert int(got[0]) == 2**63 + 1 @@ -454,49 +458,53 @@ def test_uint64_bool_mix_minmax() -> None: def test_bool_minmax_arithmetic_exact() -> None: # min of two bool expressions keeps kind 'bool', so a + a intifies - vec = vectorize(make_fn(" a = min(x > 0, x > 1)\n return a + a")) + vec = array_vectorize.vectorize(support.make_fn(" a = min(x > 0, x > 1)\n return a + a")) assert vec(np.asarray([2]))[0] == 2 def test_bool_minmax_mixed_arithmetic() -> None: - vec = vectorize(make_fn(" return max(x > 0, x > 1) + (x > 0)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x > 0, x > 1) + (x > 0)")) assert vec(np.asarray([2]))[0] == 2 def test_bool_minmax_bitwise_downstream() -> None: - vec = vectorize(make_fn(" a = max(x > 0, x > 1)\n return (a + a) & 1")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x > 0, x > 1)\n return (a + a) & 1") + ) assert vec(np.asarray([2]))[0] == 0 def test_bool_minmax_with_literals() -> None: - vec = vectorize(make_fn(" return max(x > 0, False)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x > 0, False)")) assert list(np.asarray(vec(np.asarray([2, -1])), dtype=bool)) == [True, False] def test_bool_only_minmax_bitwise_numpy() -> None: - vec = vectorize(make_fn(" return min(x > 0, x > 1) & (x > 0)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x > 0, x > 1) & (x > 0)")) # scalar: f(2) = min(True, True) & True = True; f(0) = False & False assert list(np.asarray(vec(np.asarray([2, 0])), dtype=bool)) == [True, False] def test_bool_only_minmax_or_bitwise() -> None: - vec = vectorize(make_fn(" return max(x > 0, x > 1) | (x > 5)")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x > 0, x > 1) | (x > 5)")) assert list(np.asarray(vec(np.asarray([2, 0])), dtype=bool)) == [True, False] def test_bool_int_minmax_no_int8_overflow() -> None: # min(bool, 1) selects in int64: a * 100 * 2 must not wrap int8 - vec = vectorize(make_fn(" a = min(x > 0, 1)\n return a * 100 * 2")) + vec = array_vectorize.vectorize( + support.make_fn(" a = min(x > 0, 1)\n return a * 100 * 2") + ) assert vec(np.asarray([2]))[0] == 200 def test_bool_int_minmax_large_literal() -> None: - vec = vectorize(make_fn(" a = max(x > 0, 300)\n return a * 100")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x > 0, 300)\n return a * 100")) assert vec(np.asarray([2]))[0] == 30000 def test_bool_float_minmax() -> None: - vec = vectorize(make_fn(" return min(x > 0, 1.5)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x > 0, 1.5)")) assert vec(np.asarray([2]))[0] == 1.0 @@ -506,59 +514,69 @@ def test_bool_float_minmax() -> None: def test_minmax_runtime_bool_arithmetic_numpy() -> None: # x is a boolean array at runtime, but a parameter statically: the # min/max result's bool dtype is only known at runtime - vec = vectorize(make_fn(" a = min(x, True)\n return a + a")) + vec = array_vectorize.vectorize(support.make_fn(" a = min(x, True)\n return a + a")) assert list(np.asarray(vec(np.asarray([True, False])))) == [2, 0] def test_minmax_runtime_bool_ref_chain() -> None: - vec = vectorize(make_fn(" a = min(x, True)\n b = a\n return b + b")) + vec = array_vectorize.vectorize( + support.make_fn(" a = min(x, True)\n b = a\n return b + b") + ) assert list(np.asarray(vec(np.asarray([True, False])))) == [2, 0] def test_minmax_runtime_bool_direct_call() -> None: - vec = vectorize(make_fn(" return min(x, True) + min(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, True) + min(x, True)")) assert list(np.asarray(vec(np.asarray([True, False])))) == [2, 0] def test_minmax_runtime_bool_loop_carried() -> None: - vec = vectorize( - make_fn(" a = x\n for i in range(1):\n a = min(x, True)\n return a + a") + vec = array_vectorize.vectorize( + support.make_fn( + " a = x\n for i in range(1):\n a = min(x, True)\n return a + a" + ) ) assert list(np.asarray(vec(np.asarray([True, False])))) == [2, 0] def test_minmax_runtime_numeric_noop() -> None: # the runtime intify is a no-op for numeric dtypes: min picks x (< 1) - vec = vectorize(make_fn(" a = min(x, True)\n return a + a")) + vec = array_vectorize.vectorize(support.make_fn(" a = min(x, True)\n return a + a")) assert np.allclose(vec(np.asarray([0.5, 0.0])), [1.0, 0.0]) got = vec(np.asarray([0, -2], dtype=np.int32)) assert list(np.asarray(got)) == [0, -4] def test_unary_negate_maybe_bool() -> None: - vec = vectorize(make_fn(" a = min(x, True)\n return -a")) + vec = array_vectorize.vectorize(support.make_fn(" a = min(x, True)\n return -a")) assert list(np.asarray(vec(np.asarray([True, False])))) == [-1, 0] def test_bitwise_maybe_bool_arithmetic() -> None: - vec = vectorize(make_fn(" a = min(x, True)\n b = a & True\n return b + b")) + vec = array_vectorize.vectorize( + support.make_fn(" a = min(x, True)\n b = a & True\n return b + b") + ) assert list(np.asarray(vec(np.asarray([True, False])))) == [2, 0] def test_bitwise_maybe_bool_or() -> None: - vec = vectorize(make_fn(" a = min(x, True)\n b = a | False\n return b + b")) + vec = array_vectorize.vectorize( + support.make_fn(" a = min(x, True)\n b = a | False\n return b + b") + ) assert list(np.asarray(vec(np.asarray([True, False])))) == [2, 0] def test_arithmetic_maybe_bool_result_not_maybe_bool() -> None: # a + 1 intifies inside, so the result is numeric and needs no # second conversion (values stay exact) - vec = vectorize(make_fn(" a = min(x, True)\n b = a + 1\n return b + b")) + vec = array_vectorize.vectorize( + support.make_fn(" a = min(x, True)\n b = a + 1\n return b + b") + ) assert list(np.asarray(vec(np.asarray([True, False])))) == [4, 2] def test_truediv_bool_numpy() -> None: - vec = vectorize(make_fn(" return min(x, True) / 2")) + vec = array_vectorize.vectorize(support.make_fn(" return min(x, True) / 2")) got = vec(np.asarray([True, False])) assert list(np.asarray(got)) == [0.5, 0.0] @@ -567,14 +585,14 @@ def test_truediv_bool_numpy() -> None: def test_maybe_bool_uint64_arithmetic_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a + 0")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a + 0")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) assert int(got[0]) == 9223372036854775809 def test_maybe_bool_uint64_chained_arith_no_rounding_cast() -> None: # the arith dtype must not route uint64 through float64 - vec = vectorize(make_fn(" a = max(x, True)\n return (a + 0) * 1")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return (a + 0) * 1")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) assert int(got[0]) == 9223372036854775809 @@ -583,27 +601,31 @@ def test_uint64_negative_literal_arithmetic_exact() -> None: # (2**63 + 3) + (-2) = 2**63 + 1: fits uint64; add/sub/mul keep the # modular uint64 path for negative literals, exact instead of # rounding through float64 - vec = vectorize(make_fn(" a = max(x, True)\n return a + -2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a + -2")) got = vec(np.asarray([2**63 + 3], dtype=np.uint64)) assert int(got[0]) == 9223372036854775809 def test_uint64_add_negative_literal_direct_call_exact() -> None: - vec = vectorize(make_fn(" return max(x, True) + -2")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, True) + -2")) got = vec(np.asarray([2**63 + 3], dtype=np.uint64)) assert int(got[0]) == 9223372036854775809 def test_uint64_plus_nonnegative_int64_array_exact() -> None: # huge uint64 + non-negative signed values: modular uint64 - vec = vectorize(make_fn(" a = max(x, True)\n return a + y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a + y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([0], dtype=np.int64)) assert int(got[0]) == 9223372036854775809 def test_uint64_fit_plus_negative_int64_array_exact() -> None: # uint64 values that fit int64: int64 arithmetic handles negatives - vec = vectorize(make_fn(" a = max(x, True)\n return a + y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a + y", defaults="x, y") + ) got = vec(np.asarray([5], dtype=np.uint64), np.asarray([-3], dtype=np.int64)) assert int(got[0]) == 2 @@ -611,13 +633,15 @@ def test_uint64_fit_plus_negative_int64_array_exact() -> None: def test_uint64_mixed_sign_fallback_documented() -> None: # huge uint64 + negative signed: per-lane results fit no single # dtype; float64 approximation (documented divergence) - vec = vectorize(make_fn(" a = max(x, True)\n return a + y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a + y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([-1], dtype=np.int64)) assert np.isclose(float(got[0]), 2**63) def test_uint64_negate_large_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return -a")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return -a")) got = vec(np.asarray([2**53 + 1], dtype=np.uint64)) assert int(got[0]) == -9007199254740993 @@ -625,7 +649,7 @@ def test_uint64_negate_large_exact() -> None: def test_uint64_negate_huge_best_effort_float() -> None: # values beyond int64: the result is unrepresentable exactly; the # cast falls back to float64 (documented best effort) - vec = vectorize(make_fn(" a = max(x, True)\n return -a")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return -a")) got = vec(np.asarray([2**63 + 1], dtype=np.uint64)) assert np.isclose(float(got[0]), -(2.0**63 + 1)) @@ -633,7 +657,7 @@ def test_uint64_negate_huge_best_effort_float() -> None: def test_uint64_negate_mixed_lanes_best_effort() -> None: # -(2**63 + 1) is unrepresentable in int64: the batch falls back to # float64 (documented); the small lane approximates - vec = vectorize(make_fn(" return -max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return -max(x, True)")) got = vec(np.asarray([2**53 + 1, 2**63 + 1], dtype=np.uint64)) assert np.isclose(float(np.asarray(got).reshape(-1)[0]), -9007199254740993) @@ -642,13 +666,17 @@ def test_uint64_negate_mixed_lanes_best_effort() -> None: def test_uint64_minus_signed_negative_result_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a - y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a - y", defaults="x, y") + ) got = vec(np.asarray([5], dtype=np.uint64), np.asarray([7], dtype=np.int64)) assert int(got[0]) == -2 def test_sub_nonnegative_result_keeps_uint64() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a - y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a - y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 5], dtype=np.uint64), np.asarray([3], dtype=np.int64)) assert int(got[0]) == 2**63 + 2 @@ -656,66 +684,78 @@ def test_sub_nonnegative_result_keeps_uint64() -> None: def test_sub_mixed_magnitude_fallback_documented() -> None: # per-lane differences spanning u64-only positives and int64 # negatives fit no single dtype: float64 approximation - vec = vectorize(make_fn(" a = max(x, True)\n return a - y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a - y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 5, 1], dtype=np.uint64), np.asarray([3, 4], dtype=np.int64)) assert np.isclose(float(got[0]), 2**63 + 2) assert np.isclose(float(got[1]), -3) def test_uint64_minus_uint64_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a - y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a - y", defaults="x, y") + ) got = vec(np.asarray([5], dtype=np.uint64), np.asarray([7], dtype=np.uint64)) assert int(got[0]) == -2 def test_sub_nonnegative_result_keeps_uint64_unsigned() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a - y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a - y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 5], dtype=np.uint64), np.asarray([3], dtype=np.uint64)) assert int(got[0]) == 2**63 + 2 def test_sub_mixed_magnitude_fallback_documented_unsigned() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a - y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a - y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 5, 1], dtype=np.uint64), np.asarray([3, 4], dtype=np.uint64)) assert np.isclose(float(got[0]), 2**63 + 2) assert np.isclose(float(got[1]), -3) def test_int64_minus_uint64_huge_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return y - a", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return y - a", defaults="x, y") + ) got = vec(np.asarray([2**63], dtype=np.uint64), np.asarray([2**63 - 1], dtype=np.int64)) assert int(got[0]) == -1 def test_sub_result_int64_min_exact() -> None: # -(2**63) is exactly int64-min: the mod-2**64 subtraction yields it - vec = vectorize(make_fn(" a = max(x, True)\n return y - a", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return y - a", defaults="x, y") + ) got = vec(np.asarray([2**63], dtype=np.uint64), np.asarray([0], dtype=np.int64)) assert int(got[0]) == -(2**63) def test_literal_minus_uint64_exact() -> None: - vec = vectorize(make_fn(" return 0 - max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return 0 - max(x, True)")) got = vec(np.asarray([1], dtype=np.uint64)) assert int(got[0]) == -1 def test_uint64_minus_negative_literal_exact() -> None: - vec = vectorize(make_fn(" return max(x, True) - -2")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, True) - -2")) got = vec(np.asarray([2**63 + 3], dtype=np.uint64)) assert int(got[0]) == 2**63 + 5 def test_negative_literal_minus_uint64_exact() -> None: - vec = vectorize(make_fn(" return -2 - max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return -2 - max(x, True)")) got = vec(np.asarray([5], dtype=np.uint64)) assert int(got[0]) == -7 def test_negative_literal_sub_int64_min_exact() -> None: # -(2**63 - 1) - True == -(2**63): exactly int64-min - vec = vectorize( - make_module( + vec = array_vectorize.vectorize( + support.make_module( "import math\n\n" "BOUND = -(2**63 - 1)\n\n" "def subject(x, y=2.0):\n" @@ -728,8 +768,8 @@ def test_negative_literal_sub_int64_min_exact() -> None: def test_negative_literal_sub_past_int64_min_fallback() -> None: # -(2**63) - True == -(2**63) - 1: fits no dtype, documented float64 - vec = vectorize( - make_module( + vec = array_vectorize.vectorize( + support.make_module( "import math\n\n" "BOUND = -(2**63)\n\n" "def subject(x, y=2.0):\n" @@ -745,10 +785,10 @@ def test_negative_literal_sub_past_int64_min_fallback() -> None: def test_uint64_floordiv_mod_positive_keep_uint64() -> None: # non-negative divisors keep the exact modular uint64 path - vec = vectorize(make_fn(" a = max(x, True)\n return a // 2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a // 2")) got = vec(np.asarray([2**63 + 3], dtype=np.uint64)) assert int(got[0]) == (2**63 + 3) // 2 - vec = vectorize(make_fn(" a = max(x, True)\n return a % 2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a % 2")) got = vec(np.asarray([2**63 + 3], dtype=np.uint64)) assert int(got[0]) == 1 @@ -756,21 +796,21 @@ def test_uint64_floordiv_mod_positive_keep_uint64() -> None: def test_uint64_mod_negative_divisor() -> None: # -2 % 3 == 1: a negative divisor changes mod semantics, so the # uint64 modular path must not wrap the literal - vec = vectorize(make_fn(" a = max(x, True)\n return -2 % a")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return -2 % a")) got = vec(np.asarray([3], dtype=np.uint64)) assert got[0] == 1 def test_uint64_mod_negative_divisor_formula() -> None: # 5 % -3 == -1 (not -(5 % 3) == -2); 6 % -3 == 0 - vec = vectorize(make_fn(" return max(x, True) % -3")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, True) % -3")) assert int(vec(np.asarray([5], dtype=np.uint64))[0]) == -1 assert int(vec(np.asarray([6], dtype=np.uint64))[0]) == 0 def test_uint64_mod_negative_large_exact() -> None: # (2**53 + 1) % -2 == -1: fits int64; float64 would round the value - vec = vectorize(make_fn(" a = max(x, True)\n return a % -2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a % -2")) got = vec(np.asarray([2**53 + 1], dtype=np.uint64)) assert int(got[0]) == -1 @@ -778,38 +818,38 @@ def test_uint64_mod_negative_large_exact() -> None: def test_uint64_mod_negative_mixed_lanes_exact() -> None: # both remainders fit int64: computed exactly per lane from uint64 # magnitudes, regardless of input magnitude - vec = vectorize(make_fn(" return max(x, True) % -2")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, True) % -2")) got = vec(np.asarray([2**53 + 1, 2**63 + 1], dtype=np.uint64)) assert list(map(int, np.asarray(got))) == [-1, -1] def test_uint64_mod_int64_min_divisor() -> None: - mod = make_module( + mod = support.make_module( "import math\n\n" "BOUND = -(2**63)\n\n" "def subject(x, y=2.0):\n" " return max(x, True) % BOUND\n" ) - vec = vectorize(mod.subject) + vec = array_vectorize.vectorize(mod.subject) got = vec(np.asarray([1], dtype=np.uint64)) assert int(got[0]) == -9223372036854775807 def test_uint64_floordiv_negative_divisor() -> None: # 3 // -2 == -2: floor division with a negative divisor - vec = vectorize(make_fn(" a = max(x, True)\n return a // -2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a // -2")) got = vec(np.asarray([3], dtype=np.uint64)) assert got[0] == -2.0 def test_uint64_floordiv_negative_large_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a // -2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a // -2")) got = vec(np.asarray([2**53 + 1], dtype=np.uint64)) assert int(got[0]) == -4503599627370497 def test_uint64_floordiv_negative_mixed_lanes_exact() -> None: - vec = vectorize(make_fn(" return max(x, True) // -2")) + vec = array_vectorize.vectorize(support.make_fn(" return max(x, True) // -2")) got = vec(np.asarray([2**53 + 1, 2**63 + 1], dtype=np.uint64)) assert list(map(int, np.asarray(got))) == [-4503599627370497, -4611686018427387905] @@ -817,51 +857,59 @@ def test_uint64_floordiv_negative_mixed_lanes_exact() -> None: def test_uint64_truediv_negative_divisor() -> None: # 3 / -2 == -1.5: true division is float, and the operands must be # cast to float (strict rejects integer true division of u64 wraps) - vec = vectorize(make_fn(" a = max(x, True)\n return a / -2")) + vec = array_vectorize.vectorize(support.make_fn(" a = max(x, True)\n return a / -2")) got = vec(np.asarray([3], dtype=np.uint64)) assert got[0] == -1.5 def test_uint64_negative_literal_mod_array_mixed() -> None: # literal-left remainder: (v - |d| mod v) mod v, exact in uint64 - vec = vectorize(make_fn(" return -2 % max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return -2 % max(x, True)")) got = vec(np.asarray([2**53 + 1, 2**63 + 1], dtype=np.uint64)) assert list(map(int, np.asarray(got))) == [2**53 - 1, 2**63 - 1] def test_negative_literal_mod_uint64_exact() -> None: # -1 % (2**63 + 2) == 2**63 + 1: fits uint64 (not int64) - vec = vectorize(make_fn(" return -1 % max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return -1 % max(x, True)")) got = vec(np.asarray([2**63 + 2], dtype=np.uint64)) assert int(got[0]) == 9223372036854775809 def test_uint64_negative_literal_floordiv_array() -> None: # (-d) // v == -ceil(d / v), exact in int64 - vec = vectorize(make_fn(" return -5 // max(x, True)")) + vec = array_vectorize.vectorize(support.make_fn(" return -5 // max(x, True)")) got = vec(np.asarray([2], dtype=np.uint64)) assert int(got[0]) == -3 # -5 // 2 == -3 def test_uint64_mod_positive_int64_array_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a % y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a % y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 3], dtype=np.uint64), np.asarray([2], dtype=np.int64)) assert int(got[0]) == 1 def test_uint64_floordiv_positive_int64_array_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a // y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a // y", defaults="x, y") + ) got = vec(np.asarray([2**63 + 3], dtype=np.uint64), np.asarray([2], dtype=np.int64)) assert int(got[0]) == 4611686018427387905 def test_uint64_floordiv_signed_array_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return a // y", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return a // y", defaults="x, y") + ) got = vec(np.asarray([7], dtype=np.uint64), np.asarray([-2], dtype=np.int64)) assert int(got[0]) == -4 def test_nonnegative_signed_left_mod_uint64_exact() -> None: - vec = vectorize(make_fn(" a = max(x, True)\n return y % a", defaults="x, y")) + vec = array_vectorize.vectorize( + support.make_fn(" a = max(x, True)\n return y % a", defaults="x, y") + ) got = vec(np.asarray([2**63 + 1], dtype=np.uint64), np.asarray([7], dtype=np.int64)) assert int(got[0]) == 7 diff --git a/tests/behavior/test_calls_math_casts.py b/tests/behavior/test_calls_math_casts.py index ff2250e..7211f44 100644 --- a/tests/behavior/test_calls_math_casts.py +++ b/tests/behavior/test_calls_math_casts.py @@ -3,9 +3,9 @@ from __future__ import annotations import numpy as np -from support import make_fn, vfn +import support -from array_vectorize import vectorize +import array_vectorize X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) Y = np.asarray([3.0, 1.5, 1.0, 0.5, -1.0, -3.0]) @@ -17,20 +17,20 @@ def test_math_calls() -> None: mx = X * (X > 0) # keep sqrt domain valid expected = np.sqrt(mx) + np.exp(mx) + np.log1p(mx) + np.sin(mx) - got = vfn("math_calls")(mx) + got = support.vfn("math_calls")(mx) with np.errstate(invalid="ignore"): assert np.allclose(got, expected, equal_nan=True) def test_math_two_arg() -> None: expected = np.arctan2(X, Y) + np.hypot(X, Y) + np.copysign(X, Y) - assert np.allclose(vfn("math_two_arg")(X, Y), expected) + assert np.allclose(support.vfn("math_two_arg")(X, Y), expected) def test_complex_expr() -> None: t = np.exp(-(X * X + Y * Y)) expected = t / (1.0 + t) - assert np.allclose(vfn("complex_expr")(X, Y), expected) + assert np.allclose(support.vfn("complex_expr")(X, Y), expected) # -------------------------------------------------- pow and floor-division/mod @@ -38,39 +38,41 @@ def test_complex_expr() -> None: def test_pow_op() -> None: expected = np.power(X, 2) + np.power(2.0, X) - assert np.allclose(vfn("pow_op")(X), expected) + assert np.allclose(support.vfn("pow_op")(X), expected) def test_floor_mod_ops() -> None: expected = np.floor_divide(X, 2) + np.remainder(X, 3) - got = vfn("floor_mod_ops")(X) + got = support.vfn("floor_mod_ops")(X) assert np.allclose(got, expected) def test_div_mod_sign_semantics_match_python() -> None: - fn_mod = make_fn(" return x % 3") - fn_floordiv = make_fn(" return x // 3") + fn_mod = support.make_fn(" return x % 3") + fn_floordiv = support.make_fn(" return x // 3") xs = np.asarray([-7.0, -6.5, 6.5, 7.0]) - assert np.allclose(vectorize(fn_mod)(xs), [x % 3 for x in xs]) - assert np.allclose(vectorize(fn_floordiv)(xs), [x // 3 for x in xs]) + assert np.allclose(array_vectorize.vectorize(fn_mod)(xs), [x % 3 for x in xs]) + assert np.allclose(array_vectorize.vectorize(fn_floordiv)(xs), [x // 3 for x in xs]) # --------------------------------------------------------------- abs and round def test_abs_round() -> None: - assert np.allclose(vfn("abs_round")(X), np.abs(X) + np.round(X)) + assert np.allclose(support.vfn("abs_round")(X), np.abs(X) + np.round(X)) def test_round_half_even() -> None: - fn = make_fn(" return round(x)") + fn = support.make_fn(" return round(x)") xs = np.asarray([0.5, 1.5, 2.5, -0.5, -1.5]) - assert np.array_equal(vectorize(fn)(xs), [round(x) for x in xs]) # banker's rounding + assert np.array_equal( + array_vectorize.vectorize(fn)(xs), [round(x) for x in xs] + ) # banker's rounding def test_abs_round_loop_var() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " s = x\n" " for i in range(3):\n" " s = s + abs(i - 1) + round(i / 2)\n" @@ -84,13 +86,13 @@ def test_abs_round_loop_var() -> None: def test_abs_exact_large_int() -> None: - vec = vectorize(make_fn(" return x + abs(9007199254740993)")) + vec = array_vectorize.vectorize(support.make_fn(" return x + abs(9007199254740993)")) got = vec(np.asarray([0], dtype=np.int64)) assert got[0] == 9007199254740993 def test_abs_int_bitwise() -> None: - vec = vectorize(make_fn(" return abs(3) & x")) + vec = array_vectorize.vectorize(support.make_fn(" return abs(3) & x")) got = vec(np.asarray([1], dtype=np.int64)) assert got[0] == 1 @@ -99,24 +101,24 @@ def test_abs_int_bitwise() -> None: def test_casts() -> None: - got = vfn("casts")(X) + got = support.vfn("casts")(X) expected = X.astype(np.int64) + X + X.astype(bool) + X.astype(np.int64) assert np.allclose(got, expected) def test_cast_scalar_default_param() -> None: - vec = vectorize(make_fn(" return x + int(y)")) + vec = array_vectorize.vectorize(support.make_fn(" return x + int(y)")) got = vec(np.asarray([1.0])) assert got[0] == 3.0 # y defaults to 2.0 -> int(2.0) == 2 def test_cast_literal_folds() -> None: - vec = vectorize(make_fn(" return x + int(3.0)")) + vec = array_vectorize.vectorize(support.make_fn(" return x + int(3.0)")) got = vec(np.asarray([1.0])) assert got[0] == 4.0 def test_trunc_scalar_default_param() -> None: - vec = vectorize(make_fn(" return x + math.trunc(y)")) + vec = array_vectorize.vectorize(support.make_fn(" return x + math.trunc(y)")) got = vec(np.asarray([1.0])) assert got[0] == 3.0 # math.trunc(2.0) == 2 diff --git a/tests/behavior/test_closures.py b/tests/behavior/test_closures.py index d5988a0..dcb1abe 100644 --- a/tests/behavior/test_closures.py +++ b/tests/behavior/test_closures.py @@ -4,9 +4,9 @@ import corpus import numpy as np -from support import make_module, vfn +import support -from array_vectorize import vectorize +import array_vectorize X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) @@ -15,23 +15,23 @@ def test_closure_scalar() -> None: - assert np.allclose(vfn("closure_scalar")(X), X * corpus.SCALE) + assert np.allclose(support.vfn("closure_scalar")(X), X * corpus.SCALE) def test_closure_array() -> None: - got = vfn("closure_array")(np.asarray([1.0, 1.0, 1.0])) + got = support.vfn("closure_array")(np.asarray([1.0, 1.0, 1.0])) assert np.allclose(got, np.asarray([2.0, 3.0, 4.0])) def test_kwonly_defaults() -> None: expected = X * 2.0 + 0.5 - assert np.allclose(vfn("kwonly_defaults")(X), expected) + assert np.allclose(support.vfn("kwonly_defaults")(X), expected) expected = X * 4.0 + 1.0 - assert np.allclose(vfn("kwonly_defaults")(X, scale=4.0, bias=1.0), expected) + assert np.allclose(support.vfn("kwonly_defaults")(X, scale=4.0, bias=1.0), expected) def test_hidden_params_preserve_kwonly_defaults() -> None: - mod = make_module( + mod = support.make_module( "import numpy as np\n" "C = np.asarray([2.0])\n" "\n" @@ -39,6 +39,6 @@ def test_hidden_params_preserve_kwonly_defaults() -> None: "def subject(x, *, scale=3):\n" " return x * scale + C\n" ) - vec = vectorize(mod.subject) + vec = array_vectorize.vectorize(mod.subject) assert np.allclose(vec(np.asarray([1.0])), [5.0]) assert np.allclose(vec(np.asarray([1.0]), scale=4), [6.0]) diff --git a/tests/behavior/test_control_flow.py b/tests/behavior/test_control_flow.py index d63d129..0e89b52 100644 --- a/tests/behavior/test_control_flow.py +++ b/tests/behavior/test_control_flow.py @@ -4,9 +4,9 @@ import numpy as np import pytest -from support import make_fn, vfn +import support -from array_vectorize import vectorize +import array_vectorize X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) @@ -16,49 +16,49 @@ def test_ternary() -> None: expected = np.where(X > 0, X * 2, X / 2) - assert np.allclose(vfn("ternary")(X), expected) + assert np.allclose(support.vfn("ternary")(X), expected) def test_ssa_rebind() -> None: expected = (X + 1) * 2 + 3 - assert np.allclose(vfn("ssa_rebind")(X), expected) + assert np.allclose(support.vfn("ssa_rebind")(X), expected) @pytest.mark.parametrize("name", ["relu", "relu_with_else"]) def test_relu_variants(name: str) -> None: - assert np.allclose(vfn(name)(X), np.maximum(X, 0.0)) + assert np.allclose(support.vfn(name)(X), np.maximum(X, 0.0)) def test_psi() -> None: expected = np.where(X < 0, 0.0, X * np.exp(-X)) - assert np.allclose(vfn("psi")(X), expected) + assert np.allclose(support.vfn("psi")(X), expected) def test_piecewise() -> None: expected = np.where(X < -1, -1.0, np.where(X > 1, 1.0, X)) - assert np.allclose(vfn("piecewise")(X), expected) + assert np.allclose(support.vfn("piecewise")(X), expected) def test_merge_with_prior() -> None: - assert np.allclose(vfn("merge_with_prior")(X), np.where(X > 0, X, 0.0)) + assert np.allclose(support.vfn("merge_with_prior")(X), np.where(X > 0, X, 0.0)) def test_nested_early_returns() -> None: expected = np.where(X > 0, np.where(X > 10, 1.0, 2.0), 3.0) - assert np.allclose(vfn("nested_early_returns")(X), expected) + assert np.allclose(support.vfn("nested_early_returns")(X), expected) def test_both_branches_assign() -> None: expected = np.where(X > 0, X, -X) + np.where(X > 0, 1.0, 2.0) - assert np.allclose(vfn("both_branches_assign")(X), expected) + assert np.allclose(support.vfn("both_branches_assign")(X), expected) # ----------------------------------------------- branch merging of value kinds def test_branch_merged_maybe_bool_arithmetic() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " if x == True:\n" " a = min(x, True)\n" " else:\n" @@ -73,7 +73,7 @@ def test_branch_merged_maybe_bool_arithmetic() -> None: def test_deeply_nested_boolop_chain() -> None: - vec = vectorize(make_fn(" return x and y and x and y")) + vec = array_vectorize.vectorize(support.make_fn(" return x and y and x and y")) import numpy as np xs = np.asarray([0.0, 1.0, 2.0]) diff --git a/tests/behavior/test_cse.py b/tests/behavior/test_cse.py index 855a545..5b23e3e 100644 --- a/tests/behavior/test_cse.py +++ b/tests/behavior/test_cse.py @@ -3,7 +3,7 @@ from __future__ import annotations import numpy as np -from support import just_vec, vec_of, vfn +import support X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) @@ -14,15 +14,17 @@ def test_cse_opportunity() -> None: s = np.sqrt(np.abs(X)) expected = s * s + s - assert np.allclose(vfn("cse_opportunity")(np.abs(X)), expected) + assert np.allclose(support.vfn("cse_opportunity")(np.abs(X)), expected) def test_cse_temp_does_not_collide_with_ssa_bindings() -> None: - vec, fn = vec_of(" t = x + 1\n t = t + 1\n return t + math.sin(x) + math.sin(x)") + vec, fn = support.vec_of( + " t = x + 1\n t = t + 1\n return t + math.sin(x) + math.sin(x)" + ) xs = np.asarray([1.0]) assert np.isclose(vec(xs)[0], fn(1.0)) def test_signed_zero_not_cse_merged() -> None: - vec = just_vec(" return math.copysign(x, 0.0) + math.copysign(x, -0.0)") + vec = support.just_vec(" return math.copysign(x, 0.0) + math.copysign(x, -0.0)") assert np.allclose(vec(np.asarray([1.0])), [0.0]) diff --git a/tests/behavior/test_divergences.py b/tests/behavior/test_divergences.py index e885a4c..dc3877a 100644 --- a/tests/behavior/test_divergences.py +++ b/tests/behavior/test_divergences.py @@ -3,44 +3,44 @@ from __future__ import annotations import numpy as np -from support import make_fn, vfn +import support -from array_vectorize import vectorize +import array_vectorize # --------------------------------------- numeric exceptions become IEEE values def test_divergence_div_by_zero_is_inf() -> None: with np.errstate(divide="ignore"): - got = vectorize(lambda x: 1.0 / x)(np.asarray([-1.0, 0.0, 2.0])) + got = array_vectorize.vectorize(lambda x: 1.0 / x)(np.asarray([-1.0, 0.0, 2.0])) assert np.allclose(got, [-1.0, np.inf, 0.5]) def test_divergence_neg_sqrt_is_nan() -> None: with np.errstate(invalid="ignore"): - got = vfn("cse_opportunity")(np.asarray([-4.0])) + got = support.vfn("cse_opportunity")(np.asarray([-4.0])) assert np.isnan(got[0]) def test_remainder_zero_is_nan() -> None: - fn = make_fn(" return x % 0.0") + fn = support.make_fn(" return x % 0.0") with np.errstate(invalid="ignore", divide="ignore"): - got = vectorize(fn)(np.asarray([1.0, -1.0])) + got = array_vectorize.vectorize(fn)(np.asarray([1.0, -1.0])) assert np.all(np.isnan(got)) def test_int_floordiv_zero_backend_defined() -> None: - fn = make_fn(" return x // 0") + fn = support.make_fn(" return x // 0") with np.errstate(divide="ignore", invalid="ignore"): - got = vectorize(fn)(np.asarray([3, -3], dtype=np.int64)) + got = array_vectorize.vectorize(fn)(np.asarray([3, -3], dtype=np.int64)) # NumPy: 0 with a warning; documented as backend-defined (docs/architecture.md, D3) assert list(got) == [0, 0] def test_pow_negative_base_fractional_exponent_is_nan() -> None: - fn = make_fn(" return x ** 0.5") + fn = support.make_fn(" return x ** 0.5") with np.errstate(invalid="ignore"): - got = vectorize(fn)(np.asarray([-4.0, 4.0])) + got = array_vectorize.vectorize(fn)(np.asarray([-4.0, 4.0])) assert np.isnan(got[0]) assert got[1] == 2.0 @@ -49,8 +49,8 @@ def test_pow_negative_base_fractional_exponent_is_nan() -> None: def test_and_or_eager_with_nan_lanes() -> None: - fn = make_fn(" return (x > 0) and (y > 0)") - vec = vectorize(fn) + fn = support.make_fn(" return (x > 0) and (y > 0)") + vec = array_vectorize.vectorize(fn) xs = np.asarray([1.0, -1.0, np.nan]) ys = np.asarray([1.0, 1.0, 1.0]) # NaN comparisons: NaN > 0 is False (no exception in vectorized code) @@ -61,8 +61,8 @@ def test_min_max_nan_divergence_documented() -> None: # Python: min(a, nan) returns a (nan < a is False); xp.minimum # propagates NaN. The documented mapping (left-fold xp.minimum) makes # vectorized min NaN-propagating - a documented divergence. - fn = make_fn(" return min(x, y)") - vec = vectorize(fn) + fn = support.make_fn(" return min(x, y)") + vec = array_vectorize.vectorize(fn) xs = np.asarray([1.0, 1.0]) ys = np.asarray([np.nan, 2.0]) got = vec(xs, ys) @@ -75,8 +75,8 @@ def test_where_branch_type_promotion_divergence() -> None: # branch becomes float, while Python's ternary returns the taken branch # unpromoted. Values on live lanes are unchanged, but downstream ops # can observe the type (here: int remainder-by-zero vs float NaN). - fn = make_fn(" return 0 % (x if x > 1 else x != x)") - vec = vectorize(fn) + fn = support.make_fn(" return 0 % (x if x > 1 else x != x)") + vec = array_vectorize.vectorize(fn) xs = np.asarray([0.5]) # takes the bool branch -> 0 % False # scalar Python: ZeroDivisionError; vectorized: float 0.0 % 0.0 -> NaN with np.errstate(invalid="ignore", divide="ignore"): @@ -86,7 +86,7 @@ def test_where_branch_type_promotion_divergence() -> None: def test_round_negative_zero_divergence() -> None: # Python round(-0.0) returns int 0 (sign dropped); xp.round keeps -0.0, # which atan2 can observe. Documented mapping divergence. - fn = make_fn(" return math.atan2(round(x), x)") - vec = vectorize(fn) + fn = support.make_fn(" return math.atan2(round(x), x)") + vec = array_vectorize.vectorize(fn) xs = np.asarray([-0.0]) assert vec(xs)[0] == -np.pi # scalar Python: atan2(0, -0.0) = +pi diff --git a/tests/behavior/test_docstrings.py b/tests/behavior/test_docstrings.py index a091bb8..12fcde9 100644 --- a/tests/behavior/test_docstrings.py +++ b/tests/behavior/test_docstrings.py @@ -6,8 +6,8 @@ import pytest -from array_vectorize import vectorize -from array_vectorize.codegen.docstring import extract_scalar_source +import array_vectorize +from array_vectorize.codegen import docstring def relu_fn(x: float) -> float: @@ -64,7 +64,7 @@ def marker_quoting_fn(x: float) -> float: def test_name_and_docstring() -> None: - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) assert vec.__name__ == "relu_fn" doc = inspect.getdoc(vec) or "" assert doc.splitlines()[0] == "(array-vectorized) no docstring on the scalar original." @@ -73,7 +73,7 @@ def test_name_and_docstring() -> None: def test_docstring_prefixed_summary() -> None: - vec = vectorize(sigmoid_like) + vec = array_vectorize.vectorize(sigmoid_like) doc = inspect.getdoc(vec) assert doc is not None assert doc.splitlines()[0] == "(array-vectorized) A smooth sigmoid-like squashing function." @@ -86,21 +86,21 @@ def test_docstring_prefixed_summary() -> None: def test_docstring_source_round_trip(fn) -> None: # the Notes section embeds the scalar source verbatim; cleandoc of the # generated docstring recovers it exactly - vec = vectorize(fn) + vec = array_vectorize.vectorize(fn) doc = inspect.getdoc(vec) assert doc is not None - assert extract_scalar_source(doc) == inspect.getsource(fn).strip() + assert docstring.extract_scalar_source(doc) == inspect.getsource(fn).strip() def test_docstring_numpy_style_notes() -> None: # a NumPy-style original gets a NumPy-style Notes section - vec = vectorize(numpy_doc_fn) + vec = array_vectorize.vectorize(numpy_doc_fn) doc = inspect.getdoc(vec) assert "Notes\n-----" in (doc or "") def test_docstring_existing_notes_merged() -> None: - vec = vectorize(notes_doc_fn) + vec = array_vectorize.vectorize(notes_doc_fn) doc = inspect.getdoc(vec) or "" # exactly one header line (the copy inside the embedded source is indented) assert [line for line in doc.splitlines() if line == "Notes:"] == ["Notes:"] @@ -111,5 +111,5 @@ def test_docstring_marker_quoting_degrades_to_legacy() -> None: # the original's docs quote the marker line, so the combined round # trip cannot be guaranteed: the legacy source-only docstring is # emitted instead (cleandoc-exact, like the pre-redesign format) - vec = vectorize(marker_quoting_fn) + vec = array_vectorize.vectorize(marker_quoting_fn) assert inspect.getdoc(vec) == inspect.getsource(marker_quoting_fn).strip() diff --git a/tests/behavior/test_helpers.py b/tests/behavior/test_helpers.py index 3e076a1..4c36355 100644 --- a/tests/behavior/test_helpers.py +++ b/tests/behavior/test_helpers.py @@ -4,15 +4,15 @@ import numpy as np import pytest -from support import make_fn, make_module +import support -from array_vectorize import VectorizationError, vectorize +import array_vectorize # ---------------------------------------------------------------- helper calls def test_helper_call_composition() -> None: - mod = make_module( + mod = support.make_module( "import math\n" "\n" "\n" @@ -23,19 +23,19 @@ def test_helper_call_composition() -> None: "def norm(x):\n" " return math.sqrt(square(x) + square(x))\n" ) - vec = vectorize(mod.norm) + vec = array_vectorize.vectorize(mod.norm) xs = np.asarray([3.0, 4.0]) got = vec(xs) assert np.allclose(got, np.sqrt(2 * xs**2)) assert "square_vec(" in vec.source # the helper itself is vectorized and memoized - from array_vectorize.api import _HELPER_CACHE + from array_vectorize import api - assert (mod.square, False, None) in _HELPER_CACHE + assert (mod.square, False, None) in api._HELPER_CACHE def test_helper_partial_application() -> None: - mod = make_module( + mod = support.make_module( "def scale(y, k):\n" " return y * k\n" "\n" @@ -43,12 +43,12 @@ def test_helper_partial_application() -> None: "def twice(x):\n" " return scale(x, 2.0) + scale(x, 2.0)\n" ) - vec = vectorize(mod.twice) + vec = array_vectorize.vectorize(mod.twice) assert np.allclose(vec(np.asarray([1.0, 2.0])), [4.0, 8.0]) def test_helper_with_loop_calling_helper() -> None: - mod = make_module( + mod = support.make_module( "def addmul(a, b):\n" " s = 0.0\n" " for i in range(3):\n" @@ -59,20 +59,20 @@ def test_helper_with_loop_calling_helper() -> None: "def combined(x, y):\n" " return addmul(x, y) + addmul(y, x)\n" ) - vec = vectorize(mod.combined) + vec = array_vectorize.vectorize(mod.combined) xs = np.asarray([2.0]) ys = np.asarray([3.0]) assert np.allclose(vec(xs, ys), [2 * 3 * 3 * 2]) # the loop lives in the helper's generated source, not the caller's - from array_vectorize.api import _HELPER_CACHE + from array_vectorize import api - helper_src = _HELPER_CACHE[(mod.addmul, False, None)].source + helper_src = api._HELPER_CACHE[(mod.addmul, False, None)].source assert "for i in range(0, 3):" in helper_src assert "addmul_vec" in vec.source def test_same_named_helper_does_not_shadow_caller() -> None: - fn = make_fn( + fn = support.make_fn( " return helper(x) * 10", extra=( "def make_helper():\n" @@ -82,33 +82,35 @@ def test_same_named_helper_does_not_shadow_caller() -> None: "helper = make_helper()\n" ), ) - assert vectorize(fn)(np.asarray([1.0]))[0] == 20.0 + assert array_vectorize.vectorize(fn)(np.asarray([1.0]))[0] == 20.0 def test_helper_cache_protect_flag() -> None: - mod = make_module( + mod = support.make_module( "import math\n\n\n" "def inner(x):\n" " return math.sqrt(x) if x > 0 else 0.0\n\n\n" "def outer(x):\n" " return inner(x)\n" ) - vectorize(mod.outer) # caches the unprotected inner - protected = vectorize(mod.outer, protect_domains=True) + array_vectorize.vectorize(mod.outer) # caches the unprotected inner + protected = array_vectorize.vectorize(mod.outer, protect_domains=True) with np.errstate(all="raise"): got = protected(np.asarray([-1.0, 4.0])) assert np.allclose(got, [0.0, 2.0]) def test_recursive_helper_rejected() -> None: - mod = make_module( + mod = support.make_module( "def fib(n):\n if n < 2:\n return n\n return fib(n - 1) + fib(n - 2)\n" ) - with pytest.raises(VectorizationError, match="recursive"): - vectorize(mod.fib) + with pytest.raises(array_vectorize.VectorizationError, match="recursive"): + array_vectorize.vectorize(mod.fib) def test_helper_keyword_call_rejected() -> None: - fn = make_fn(" return helper(x, y=9)", extra="def helper(x, y=2):\n return x + y\n") - with pytest.raises(VectorizationError, match="keyword"): - vectorize(fn) + fn = support.make_fn( + " return helper(x, y=9)", extra="def helper(x, y=2):\n return x + y\n" + ) + with pytest.raises(array_vectorize.VectorizationError, match="keyword"): + array_vectorize.vectorize(fn) diff --git a/tests/behavior/test_inspectability.py b/tests/behavior/test_inspectability.py index d53d5bb..d4c1c75 100644 --- a/tests/behavior/test_inspectability.py +++ b/tests/behavior/test_inspectability.py @@ -9,7 +9,7 @@ import pytest -from array_vectorize import VectorizationError, get_source, vectorize +import array_vectorize def relu_fn(x: float) -> float: @@ -29,13 +29,13 @@ def sigmoid_like(x: float, alpha: float = 1.0) -> float: def test_source_attribute() -> None: - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) assert isinstance(vec.source, str) assert "xp.where" in vec.source def test_getsource_returns_generated_not_original() -> None: - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) src = inspect.getsource(vec) assert "array_namespace" in src assert "def relu_fn_vec" in src @@ -43,14 +43,14 @@ def test_getsource_returns_generated_not_original() -> None: def test_getsourcelines() -> None: - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) lines, lineno = inspect.getsourcelines(vec) assert any("array_namespace" in line for line in lines) assert lineno >= 1 def test_linecache_survives_checkcache() -> None: - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) filename = vec.__code__.co_filename assert filename == "" linecache.checkcache(filename) @@ -59,7 +59,7 @@ def test_linecache_survives_checkcache() -> None: def test_traceback_renders_generated_lines() -> None: - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) try: vec(1.0) # no array among args -> namespace resolution fails except TypeError: @@ -70,7 +70,7 @@ def test_traceback_renders_generated_lines() -> None: def test_wrapped_points_to_original() -> None: - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) # __wrapped__ is deliberately not set (it would hide the generated source # from inspect.getsourcelines, which unwraps unconditionally); the # original is reachable via the marker and drives the signature. @@ -81,7 +81,7 @@ def test_wrapped_points_to_original() -> None: def test_module_and_qualname_from_original() -> None: # module attribution makes pydoc render "Help on function f in module m" # and is safe: getsource/getsourcelines use co_filename + linecache - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) assert vec.__module__ == relu_fn.__module__ assert vec.__qualname__ == relu_fn.__qualname__ @@ -90,7 +90,7 @@ def test_pickle_fails_loudly_never_aliases() -> None: # __module__/__qualname__ match the original, so pickle resolves the # name and fails the identity check — it cannot silently serialize # (and unpickle to) the scalar original - vec = vectorize(relu_fn) + vec = array_vectorize.vectorize(relu_fn) with pytest.raises(pickle.PicklingError): pickle.dumps(vec) @@ -99,7 +99,7 @@ def test_signature_preserved() -> None: def with_defaults(x: float, scale: float = 2.0, *, bias: float = 0.5) -> float: return x * scale + bias - vec = vectorize(with_defaults) + vec = array_vectorize.vectorize(with_defaults) sig = inspect.signature(vec) assert list(sig.parameters) == ["x", "scale", "bias"] assert sig.parameters["scale"].default == 2.0 @@ -108,21 +108,21 @@ def with_defaults(x: float, scale: float = 2.0, *, bias: float = 0.5) -> float: def test_get_source_helper() -> None: - vec = vectorize(relu_fn) - assert get_source(vec) == vec.source - with pytest.raises(VectorizationError, match="not a vectorized function"): - get_source(relu_fn) + vec = array_vectorize.vectorize(relu_fn) + assert array_vectorize.get_source(vec) == vec.source + with pytest.raises(array_vectorize.VectorizationError, match="not a vectorized function"): + array_vectorize.get_source(relu_fn) def test_help_renders() -> None: # help() must not raise; it shows the module header, the prefixed # summary, and the original source in the Notes section - vec = vectorize(sigmoid_like) + vec = array_vectorize.vectorize(sigmoid_like) + import contextlib import io - from contextlib import redirect_stdout buf = io.StringIO() - with redirect_stdout(buf): + with contextlib.redirect_stdout(buf): help(vec) out = buf.getvalue() assert f"Help on function sigmoid_like in module {sigmoid_like.__module__}:" in out diff --git a/tests/behavior/test_lambdas.py b/tests/behavior/test_lambdas.py index b628fa4..5676d70 100644 --- a/tests/behavior/test_lambdas.py +++ b/tests/behavior/test_lambdas.py @@ -5,13 +5,13 @@ import math import numpy as np -from support import make_module +import support -from array_vectorize import vectorize +import array_vectorize def test_lambda_vectorize() -> None: - vec = vectorize(lambda x: x * 2.0) + vec = array_vectorize.vectorize(lambda x: x * 2.0) got = vec(np.asarray([1.0, 2.0])) assert np.allclose(got, [2.0, 4.0]) assert "def lambda_vec" in vec.source @@ -21,36 +21,40 @@ def test_lambda_vectorize() -> None: def test_multiple_lambdas_on_one_line() -> None: - mod = make_module("f1, f2 = (lambda x: x + 1), (lambda x: x + 2)\n") - assert vectorize(mod.f2)(np.asarray([1.0]))[0] == 3.0 - assert vectorize(mod.f1)(np.asarray([1.0]))[0] == 2.0 + mod = support.make_module("f1, f2 = (lambda x: x + 1), (lambda x: x + 2)\n") + assert array_vectorize.vectorize(mod.f2)(np.asarray([1.0]))[0] == 3.0 + assert array_vectorize.vectorize(mod.f1)(np.asarray([1.0]))[0] == 2.0 def test_lambdas_differing_only_in_names() -> None: - mod = make_module("import math\nf1, f2 = (lambda x: math.sin(x)), (lambda x: math.cos(x))\n") - assert np.isclose(vectorize(mod.f2)(np.asarray([1.0]))[0], math.cos(1.0)) - assert np.isclose(vectorize(mod.f1)(np.asarray([1.0]))[0], math.sin(1.0)) + mod = support.make_module( + "import math\nf1, f2 = (lambda x: math.sin(x)), (lambda x: math.cos(x))\n" + ) + assert np.isclose(array_vectorize.vectorize(mod.f2)(np.asarray([1.0]))[0], math.cos(1.0)) + assert np.isclose(array_vectorize.vectorize(mod.f1)(np.asarray([1.0]))[0], math.sin(1.0)) def test_lambdas_differing_only_in_defaults() -> None: - mod = make_module("g1, g2 = (lambda x, y=1: x + y), (lambda x, y=2: x + y)\n") - assert vectorize(mod.g2)(np.asarray([1.0]))[0] == 3.0 - assert vectorize(mod.g1)(np.asarray([1.0]))[0] == 2.0 + mod = support.make_module("g1, g2 = (lambda x, y=1: x + y), (lambda x, y=2: x + y)\n") + assert array_vectorize.vectorize(mod.g2)(np.asarray([1.0]))[0] == 3.0 + assert array_vectorize.vectorize(mod.g1)(np.asarray([1.0]))[0] == 2.0 def test_lambda_signed_zero_constants() -> None: - mod = make_module( + mod = support.make_module( "import math\n" "f1, f2 = (\n" " lambda x: math.copysign(x, 0.0),\n" " lambda x: math.copysign(x, -0.0),\n" ")\n" ) - assert vectorize(mod.f2)(np.asarray([1.0]))[0] == -1.0 - assert vectorize(mod.f1)(np.asarray([1.0]))[0] == 1.0 + assert array_vectorize.vectorize(mod.f2)(np.asarray([1.0]))[0] == -1.0 + assert array_vectorize.vectorize(mod.f1)(np.asarray([1.0]))[0] == 1.0 def test_lambda_int_vs_bool_vs_float_constants() -> None: - mod = make_module("g1, g2, g3 = (lambda x: x + 1), (lambda x: x + True), (lambda x: x + 1.0)\n") - assert vectorize(mod.g2)(np.asarray([1.0]))[0] == 2.0 - assert vectorize(mod.g3)(np.asarray([1.0]))[0] == 2.0 + mod = support.make_module( + "g1, g2, g3 = (lambda x: x + 1), (lambda x: x + True), (lambda x: x + 1.0)\n" + ) + assert array_vectorize.vectorize(mod.g2)(np.asarray([1.0]))[0] == 2.0 + assert array_vectorize.vectorize(mod.g3)(np.asarray([1.0]))[0] == 2.0 diff --git a/tests/behavior/test_loops.py b/tests/behavior/test_loops.py index fb74ca5..5d43c52 100644 --- a/tests/behavior/test_loops.py +++ b/tests/behavior/test_loops.py @@ -4,16 +4,18 @@ import numpy as np import pytest -from support import just_vec, make_fn, vec_of +import support -from array_vectorize import VectorizationError, vectorize +import array_vectorize # ------------------------------------------------------ range forms and bounds def test_loop_with_start_stop_step() -> None: - vec = vectorize( - make_fn(" s = 0.0\n for i in range(1, 5, 2):\n s = s + i * x\n return s") + vec = array_vectorize.vectorize( + support.make_fn( + " s = 0.0\n for i in range(1, 5, 2):\n s = s + i * x\n return s" + ) ) got = vec(np.asarray([1.0])) assert np.allclose(got, [(1 + 3) * 1.0]) @@ -21,23 +23,25 @@ def test_loop_with_start_stop_step() -> None: def test_loop_negative_step() -> None: - vec = vectorize( - make_fn(" s = 0.0\n for i in range(3, 0, -1):\n s = s + i * x\n return s") + vec = array_vectorize.vectorize( + support.make_fn( + " s = 0.0\n for i in range(3, 0, -1):\n s = s + i * x\n return s" + ) ) assert np.allclose(vec(np.asarray([1.0])), [6.0]) assert "for i in range(3, 0, -1):" in vec.source def test_range_two_args() -> None: - vec = vectorize( - make_fn(" s = x\n for i in range(2, 5):\n s = s + i\n return s") + vec = array_vectorize.vectorize( + support.make_fn(" s = x\n for i in range(2, 5):\n s = s + i\n return s") ) assert vec(np.asarray([0.0]))[0] == 9.0 def test_loop_bound_closure_scalar() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " s = 0.0\n for i in range(N):\n s = s + x\n return s", extra="N = 4" ) ) @@ -45,8 +49,8 @@ def test_loop_bound_closure_scalar() -> None: def test_loop_bound_expression_folds() -> None: - vec = vectorize( - make_fn(" s = 0.0\n for i in range(1 + 2):\n s = s + x\n return s") + vec = array_vectorize.vectorize( + support.make_fn(" s = 0.0\n for i in range(1 + 2):\n s = s + x\n return s") ) assert "for i in range(0, 3):" in vec.source @@ -55,7 +59,9 @@ def test_loop_bound_expression_folds() -> None: def test_simple_accumulation() -> None: - vec = vectorize(make_fn(" s = 0.0\n for i in range(3):\n s = s + x\n return s")) + vec = array_vectorize.vectorize( + support.make_fn(" s = 0.0\n for i in range(3):\n s = s + x\n return s") + ) got = vec(np.asarray([1.0, 2.0])) assert np.allclose(got, [3.0, 6.0]) assert "for i in range(0, 3):" in vec.source @@ -63,15 +69,15 @@ def test_simple_accumulation() -> None: def test_loop_index_used() -> None: - vec = vectorize( - make_fn(" s = 0.0\n for i in range(4):\n s = s + i\n return s + x") + vec = array_vectorize.vectorize( + support.make_fn(" s = 0.0\n for i in range(4):\n s = s + i\n return s + x") ) assert np.allclose(vec(np.asarray([0.0, 1.0])), [6.0, 7.0]) def test_loop_carried_two_vars() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " a = 1.0\n" " b = 0.0\n" " for i in range(3):\n" @@ -85,8 +91,8 @@ def test_loop_carried_two_vars() -> None: def test_loop_var_shadow_param() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " n = 0.0\n for n in range(3):\n n2 = n\n return n2 + x", defaults="x, n=3", ) @@ -95,8 +101,8 @@ def test_loop_var_shadow_param() -> None: def test_loop_over_closure_array() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " s = 0.0\n for i in range(2):\n s = s + x\n return s + GLOBAL_ARR", extra="GLOBAL_ARR = np.asarray([2.0, 4.0])", ) @@ -106,8 +112,8 @@ def test_loop_over_closure_array() -> None: def test_branch_inside_loop() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " s = 0.0\n" " for i in range(4):\n" " if x > 0:\n" @@ -123,8 +129,10 @@ def test_branch_inside_loop() -> None: def test_loop_differential() -> None: - fn = make_fn(" s = 1.0\n for i in range(5):\n s = s * (1.0 + x)\n return s") - vec = vectorize(fn) + fn = support.make_fn( + " s = 1.0\n for i in range(5):\n s = s * (1.0 + x)\n return s" + ) + vec = array_vectorize.vectorize(fn) for x in (-0.2, 0.0, 0.3, 1.0): expected = 1.0 for _ in range(5): @@ -133,8 +141,8 @@ def test_loop_differential() -> None: def test_loop_with_protect_domains_matches_numpy() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " a = x\n for i in range(2):\n a = a + math.sqrt(abs(x) + 1)\n return a" ), protect_domains=True, @@ -150,12 +158,14 @@ def test_loop_with_protect_domains_matches_numpy() -> None: def test_loop_local_var() -> None: - vec = vectorize(make_fn(" for i in range(3):\n t = i * x\n return t + x")) + vec = array_vectorize.vectorize( + support.make_fn(" for i in range(3):\n t = i * x\n return t + x") + ) assert np.allclose(vec(np.asarray([2.0])), [4.0 + 2.0]) def test_empty_loop_body_emits_pass() -> None: - vec = just_vec(" for i in range(3):\n pass\n return x + i") + vec = support.just_vec(" for i in range(3):\n pass\n return x + i") assert np.allclose(vec(np.asarray([1.0])), [3.0]) @@ -167,7 +177,7 @@ def test_loop_dce_drops_dead_body_bindings() -> None: " s = s + x\n" " return s" ) - vec = vectorize(make_fn(body)) + vec = array_vectorize.vectorize(support.make_fn(body)) # the original source in the docstring mentions 'dead' once; the code must not assert vec.source.count("dead") == 1 @@ -181,7 +191,7 @@ def test_loop_dce_keeps_cross_iteration_dependencies() -> None: " b = b + 1\n" " return a" ) - vec, fn = vec_of(body) + vec, fn = support.vec_of(body) xs = np.asarray([1.0, 2.0]) assert np.allclose(vec(xs), [fn(float(v)) for v in xs]) @@ -190,20 +200,24 @@ def test_loop_dce_keeps_cross_iteration_dependencies() -> None: def test_zero_trip_loop_matches_python() -> None: - vec = vectorize(make_fn(" s = 5.0\n for i in range(0):\n s = s + x\n return s")) + vec = array_vectorize.vectorize( + support.make_fn(" s = 5.0\n for i in range(0):\n s = s + x\n return s") + ) assert np.allclose(vec(np.asarray([1.0])), [5.0]) def test_zero_trip_loop_keeps_loop_var() -> None: - vec = just_vec( + vec = support.just_vec( " i = x\n a = x\n for i in range(0):\n a = a + 1\n return i + a" ) assert vec(np.asarray([1.0]))[0] == 2.0 def test_zero_trip_loop_literal_not_substituted() -> None: - vec = vectorize( - make_fn(" a = x\n for i in range(0):\n a = 4.0\n return math.sqrt(a)") + vec = array_vectorize.vectorize( + support.make_fn( + " a = x\n for i in range(0):\n a = 4.0\n return math.sqrt(a)" + ) ) assert np.allclose(vec(np.asarray([9.0, 16.0])), [3.0, 4.0]) @@ -212,15 +226,17 @@ def test_zero_trip_loop_literal_not_substituted() -> None: def test_executed_loop_literal_not_substituted() -> None: - vec = vectorize( - make_fn(" a = 4.0\n for i in range(1):\n a = 9.0\n return math.sqrt(a)") + vec = array_vectorize.vectorize( + support.make_fn( + " a = 4.0\n for i in range(1):\n a = 9.0\n return math.sqrt(a)" + ) ) assert np.allclose(vec(np.asarray([9.0])), [3.0]) def test_conditional_loop_literal_not_substituted() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " a = 4.0\n" " for i in range(1):\n" " if x > 0:\n" @@ -237,8 +253,8 @@ def test_conditional_loop_literal_not_substituted() -> None: def test_nested_loops() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " s = 0.0\n" " for i in range(3):\n" " for j in range(2):\n" @@ -250,8 +266,8 @@ def test_nested_loops() -> None: def test_nested_loop_target_carried() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " j = x\n" " s = 0.0\n" " for i in range(2):\n" @@ -265,8 +281,8 @@ def test_nested_loop_target_carried() -> None: def test_nested_loop_same_var_name() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " for i in range(2):\n for i in range(3):\n pass\n return x + i" ) ) @@ -277,7 +293,7 @@ def test_nested_loop_same_var_name() -> None: def test_loop_carried_kind_mixing() -> None: - vec, fn = vec_of( + vec, fn = support.vec_of( " b = x > 0\n" " a = x\n" " for i in range(2):\n" @@ -290,15 +306,17 @@ def test_loop_carried_kind_mixing() -> None: def test_zero_trip_loop_kind_mixing() -> None: - vec = vectorize( - make_fn(" b = x > 0\n for i in range(0):\n b = x + 0.5\n return b + b") + vec = array_vectorize.vectorize( + support.make_fn( + " b = x > 0\n for i in range(0):\n b = x + 0.5\n return b + b" + ) ) assert list(map(int, vec(np.asarray([1.0])))) == [2] def test_int_to_bool_loop_transition() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " b = 1\n" " a = x\n" " for i in range(2):\n" @@ -307,13 +325,13 @@ def test_int_to_bool_loop_transition() -> None: " return a" ) ) - assert vectorize(make_fn(" return 0")) is not None + assert array_vectorize.vectorize(support.make_fn(" return 0")) is not None assert vec(np.asarray([1.0]))[0] == 2 def test_float_to_bool_loop_transition() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " b = 1.5\n" " a = x\n" " for i in range(2):\n" @@ -327,8 +345,8 @@ def test_float_to_bool_loop_transition() -> None: def test_chained_loop_kind_mixing() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " a = 1\n" " b = 1\n" " c = 1\n" @@ -347,8 +365,8 @@ def test_chained_loop_kind_mixing() -> None: def test_float_bool_loop_mix_keeps_float_values() -> None: # b is bool in iteration 1, float afterwards: the intify must not # truncate the float iterations (float64 for float-mixed names) - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " b = x > 0\n" " a = x\n" " for i in range(2):\n" @@ -361,8 +379,8 @@ def test_float_bool_loop_mix_keeps_float_values() -> None: def test_float_mixed_propagates_through_assignment() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " b = x\n for i in range(1):\n b = x + 0.5\n c = b\n return c + c" ) ) @@ -370,8 +388,8 @@ def test_float_mixed_propagates_through_assignment() -> None: def test_true_division_loop_mix_keeps_float() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " b = x > 0\n" " a = x\n" " for i in range(2):\n" @@ -384,22 +402,24 @@ def test_true_division_loop_mix_keeps_float() -> None: def test_mixed_loop_operand_not_truncated_by_literal() -> None: - vec = vectorize( - make_fn(" b = x\n for i in range(1):\n b = x + 0.5\n return b + 1") + vec = array_vectorize.vectorize( + support.make_fn(" b = x\n for i in range(1):\n b = x + 0.5\n return b + 1") ) assert np.allclose(vec(np.asarray([1.0, 2.0])), [2.5, 3.5]) def test_mixed_loop_operand_unary_negate() -> None: - vec = vectorize( - make_fn(" b = x\n for i in range(1):\n b = x + 0.5\n return -b") + vec = array_vectorize.vectorize( + support.make_fn(" b = x\n for i in range(1):\n b = x + 0.5\n return -b") ) assert np.allclose(vec(np.asarray([1.0, 2.0])), [-1.5, -2.5]) def test_bool_loop_count_128() -> None: - vec = vectorize( - make_fn(" s = 0\n for i in range(128):\n s = s + (x > 0)\n return s") + vec = array_vectorize.vectorize( + support.make_fn( + " s = 0\n for i in range(128):\n s = s + (x > 0)\n return s" + ) ) assert vec(np.asarray([1.0]))[0] == 128 @@ -408,9 +428,9 @@ def test_bool_loop_count_128() -> None: def test_loop_bound_param_rejected() -> None: - with pytest.raises(VectorizationError, match="constant ints"): - vectorize( - make_fn( + with pytest.raises(array_vectorize.VectorizationError, match="constant ints"): + array_vectorize.vectorize( + support.make_fn( " s = 0.0\n for i in range(n):\n s = s + x\n return s", defaults="x, n=3", ) @@ -418,19 +438,23 @@ def test_loop_bound_param_rejected() -> None: def test_loop_step_zero_rejected() -> None: - with pytest.raises(VectorizationError, match="step cannot be zero"): - vectorize(make_fn(" for i in range(3, 0, 0):\n pass\n return x")) + with pytest.raises(array_vectorize.VectorizationError, match="step cannot be zero"): + array_vectorize.vectorize( + support.make_fn(" for i in range(3, 0, 0):\n pass\n return x") + ) def test_return_inside_loop_rejected() -> None: - with pytest.raises(VectorizationError, match="returns inside loop"): - vectorize( - make_fn( + with pytest.raises(array_vectorize.VectorizationError, match="returns inside loop"): + array_vectorize.vectorize( + support.make_fn( " for i in range(3):\n if x > 0:\n return 1.0\n return 0.0" ) ) def test_assign_loop_var_rejected() -> None: - with pytest.raises(VectorizationError, match="loop variable"): - vectorize(make_fn(" for i in range(3):\n i = i + 1\n return x")) + with pytest.raises(array_vectorize.VectorizationError, match="loop variable"): + array_vectorize.vectorize( + support.make_fn(" for i in range(3):\n i = i + 1\n return x") + ) diff --git a/tests/behavior/test_memoization.py b/tests/behavior/test_memoization.py index fa957e3..a4b1cbb 100644 --- a/tests/behavior/test_memoization.py +++ b/tests/behavior/test_memoization.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from array_vectorize import vectorize +import array_vectorize def relu_fn(x: float) -> float: @@ -33,31 +33,31 @@ def helper_outer_fn(x: float) -> float: def test_backref_on_canonical_vectorization() -> None: - vec = vectorize(sigmoid_like) + vec = array_vectorize.vectorize(sigmoid_like) assert sigmoid_like.__array_vectorized__ is vec def test_vectorize_memoized_per_original() -> None: - assert vectorize(sigmoid_like) is vectorize(sigmoid_like) - assert vectorize(sigmoid_like) is sigmoid_like.__array_vectorized__ + assert array_vectorize.vectorize(sigmoid_like) is array_vectorize.vectorize(sigmoid_like) + assert array_vectorize.vectorize(sigmoid_like) is sigmoid_like.__array_vectorized__ def test_revectorization_returns_same_object() -> None: - vec = vectorize(sigmoid_like) - assert vectorize(vec) is vec + vec = array_vectorize.vectorize(sigmoid_like) + assert array_vectorize.vectorize(vec) is vec def test_verify_runs_on_cache_hit() -> None: args = (np.asarray([-1.0, 0.5, 2.0]),) - vectorize(sigmoid_like, verify=args) - vectorize(sigmoid_like, verify=args) + array_vectorize.vectorize(sigmoid_like, verify=args) + array_vectorize.vectorize(sigmoid_like, verify=args) def test_backref_not_overwritten_by_variants() -> None: - vec = vectorize(sigmoid_like) - vectorize(sigmoid_like, protect_domains=True) + vec = array_vectorize.vectorize(sigmoid_like) + array_vectorize.vectorize(sigmoid_like, protect_domains=True) assert sigmoid_like.__array_vectorized__ is vec - vectorize(sigmoid_like, namespace=np) + array_vectorize.vectorize(sigmoid_like, namespace=np) assert sigmoid_like.__array_vectorized__ is vec @@ -68,21 +68,21 @@ def bad(x: float) -> float: return x with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(bad, fallback=True) + vec = array_vectorize.vectorize(bad, fallback=True) assert bad.__array_vectorized__ is vec def test_helpers_do_not_get_backref() -> None: # D7 helper compilations are internal: the user's helper function # must not grow an __array_vectorized__ marker - vectorize(helper_outer_fn) + array_vectorize.vectorize(helper_outer_fn) assert not hasattr(helper_inner_fn, "__array_vectorized__") assert helper_outer_fn.__array_vectorized__ is not None def test_idempotent_revectorization() -> None: - vec1 = vectorize(relu_fn) - vec2 = vectorize(vec1) + vec1 = array_vectorize.vectorize(relu_fn) + vec2 = array_vectorize.vectorize(vec1) x = np.asarray([-1.0, 2.0]) assert np.allclose(vec1(x), vec2(x)) assert "xp.where" in vec2.source diff --git a/tests/behavior/test_name_allocation.py b/tests/behavior/test_name_allocation.py index 2ffef55..5526fe7 100644 --- a/tests/behavior/test_name_allocation.py +++ b/tests/behavior/test_name_allocation.py @@ -5,9 +5,9 @@ import math import numpy as np -from support import make_fn, make_module, vfn +import support -from array_vectorize import vectorize +import array_vectorize X = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0, 10.0]) @@ -16,41 +16,43 @@ def test_reserved_param() -> None: - assert np.allclose(vfn("reserved_param")(X), X + 1) + assert np.allclose(support.vfn("reserved_param")(X), X + 1) def test_reserved_xp_param_keyword_call() -> None: - mod = make_module("import numpy as np\n\n\ndef subject(xp):\n return xp + 1\n") - vec = vectorize(mod.subject) + mod = support.make_module("import numpy as np\n\n\ndef subject(xp):\n return xp + 1\n") + vec = array_vectorize.vectorize(mod.subject) assert np.allclose(vec(xp=np.asarray([1.0])), [2.0]) assert "def subject_vec(xp):" in vec.source def test_namespace_var_rebinding_no_collision() -> None: - mod = make_module( + mod = support.make_module( "import math\n\n\ndef subject(xp):\n xp = xp + 1\n return math.sin(xp)\n" ) - vec = vectorize(mod.subject) + vec = array_vectorize.vectorize(mod.subject) assert np.allclose(vec(xp=np.asarray([1.0])), [math.sin(2.0)]) def test_reserved_caller_name_param_not_renamed() -> None: - mod = make_module( + mod = support.make_module( "def reserved_caller(reserved_caller_vec):\n return reserved_caller_vec + 1\n" ) - vec = vectorize(mod.reserved_caller) + vec = array_vectorize.vectorize(mod.reserved_caller) assert np.allclose(vec(reserved_caller_vec=np.asarray([1.0])), [2.0]) def test_helper_name_shadowed_by_param() -> None: - vec = vectorize(make_fn(" return min(_vec_common_dtype, 3)", defaults="_vec_common_dtype")) + vec = array_vectorize.vectorize( + support.make_fn(" return min(_vec_common_dtype, 3)", defaults="_vec_common_dtype") + ) got = vec(np.asarray([1.0])) assert got[0] == 1.0 def test_helper_names_shadowed_by_params() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " return min(_vec_arith_dtype, _vec_common_dtype) + (_vec_arith_dtype > 0)", defaults="_vec_arith_dtype, _vec_common_dtype=2.0", ) diff --git a/tests/behavior/test_public_api.py b/tests/behavior/test_public_api.py index eb95913..b410cd6 100644 --- a/tests/behavior/test_public_api.py +++ b/tests/behavior/test_public_api.py @@ -6,7 +6,6 @@ import pytest import array_vectorize -from array_vectorize import VectorizationError, vectorize def test_package_exposes_version() -> None: @@ -20,5 +19,5 @@ def test_unhashable_input_rejected_with_vectorization_error(bad: object) -> None # the canonical cache lookup runs before extraction and must not # surface a TypeError for unhashable junk: the extractor's positioned # VectorizationError ("expected a plain Python function, got ...") wins - with pytest.raises(VectorizationError, match="plain Python function"): - vectorize(bad) # type: ignore[arg-type] + with pytest.raises(array_vectorize.VectorizationError, match="plain Python function"): + array_vectorize.vectorize(bad) # type: ignore[arg-type] diff --git a/tests/conftest.py b/tests/conftest.py index b9b8b02..b5874ad 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,20 +2,20 @@ from __future__ import annotations +import importlib import os -from importlib import import_module -from types import ModuleType +import types +import hypothesis import pytest -from hypothesis import settings # The only place Hypothesis profiles are set: CI=1 runs 300 derandomized # examples per property; local runs keep Hypothesis' defaults. Neither has a # deadline: properties check values, and the first example of a test pays # the one-time vectorize() compile, which can exceed the 200 ms default. -settings.register_profile("ci", derandomize=True, max_examples=300, deadline=None) -settings.register_profile("dev", deadline=None) -settings.load_profile("ci" if os.environ.get("CI") else "dev") +hypothesis.settings.register_profile("ci", derandomize=True, max_examples=300, deadline=None) +hypothesis.settings.register_profile("dev", deadline=None) +hypothesis.settings.load_profile("ci" if os.environ.get("CI") else "dev") def pytest_addoption(parser: pytest.Parser) -> None: @@ -29,12 +29,12 @@ def pytest_addoption(parser: pytest.Parser) -> None: @pytest.fixture -def xps() -> ModuleType: +def xps() -> types.ModuleType: """The array-api-strict backend (lazy: non-strict runs skip the import).""" - return import_module("array_api_strict") + return importlib.import_module("array_api_strict") @pytest.fixture(scope="session") -def corpus() -> ModuleType: +def corpus() -> types.ModuleType: """The shared scalar-function corpus (tests/corpus.py, via pythonpath).""" - return import_module("corpus") + return importlib.import_module("corpus") diff --git a/tests/corpus.py b/tests/corpus.py index dbde376..26dc042 100644 --- a/tests/corpus.py +++ b/tests/corpus.py @@ -9,8 +9,6 @@ import numpy as np -from array_vectorize import vectorize - SCALE = 3.0 ARR = np.asarray([1.0, 2.0, 3.0]) @@ -342,4 +340,4 @@ def docstring_numpy_notes(x): "complex_expr", ] -__all__ = ["GOLDEN_NAMES", "vectorize"] +__all__ = ["GOLDEN_NAMES"] diff --git a/tests/differential/test_scalar_equivalence.py b/tests/differential/test_scalar_equivalence.py index 860103e..fdf5647 100644 --- a/tests/differential/test_scalar_equivalence.py +++ b/tests/differential/test_scalar_equivalence.py @@ -5,11 +5,11 @@ from typing import Any import corpus +import hypothesis import numpy as np -from hypothesis import given from hypothesis import strategies as st -from array_vectorize import vectorize +import array_vectorize CORPUS = corpus @@ -28,26 +28,26 @@ def arrays( def check(fn: Any, xs: np.ndarray, *rest: np.ndarray) -> None: - vec = vectorize(fn) + vec = array_vectorize.vectorize(fn) got = vec(xs, *rest) pairs = list(zip(xs, *rest, strict=True)) if rest else [(x,) for x in xs] expected = np.asarray([fn(*pair) for pair in pairs]) assert np.allclose(got, expected, equal_nan=True, rtol=1e-12), (got, expected) -@given(arrays(finite)) +@hypothesis.given(arrays(finite)) def test_relu_diff(xs: np.ndarray) -> None: check(CORPUS.relu, xs) -@given(arrays(finite)) +@hypothesis.given(arrays(finite)) def test_psi_diff(xs: np.ndarray) -> None: check(CORPUS.psi, xs) -@given(arrays(small), arrays(small)) +@hypothesis.given(arrays(small), arrays(small)) def test_clamp_diff(xs: np.ndarray, lo: np.ndarray) -> None: - vec = vectorize(CORPUS.clamp) + vec = array_vectorize.vectorize(CORPUS.clamp) lo_arr = np.minimum(np.maximum(lo, -2.0), 0.5) hi_arr = lo_arr + 1.0 got = vec(xs, lo_arr, hi_arr) @@ -57,12 +57,12 @@ def test_clamp_diff(xs: np.ndarray, lo: np.ndarray) -> None: assert np.allclose(got, expected, equal_nan=True) -@given(arrays(finite)) +@hypothesis.given(arrays(finite)) def test_piecewise_diff(xs: np.ndarray) -> None: check(CORPUS.piecewise, xs) -@given(arrays(finite)) +@hypothesis.given(arrays(finite)) def test_ternary_diff(xs: np.ndarray) -> None: check(CORPUS.ternary, xs) @@ -77,91 +77,91 @@ def test_ternary_diff(xs: np.ndarray) -> None: ) -@given(nonsingular_lanes) +@hypothesis.given(nonsingular_lanes) def test_arith_diff(lanes: list[tuple[float, float]]) -> None: xs, ys = (np.asarray(col, dtype=np.float64) for col in zip(*lanes, strict=True)) check(CORPUS.arith_ops, xs, ys) -@given(arrays(small), arrays(small)) +@hypothesis.given(arrays(small), arrays(small)) def test_boolop_numeric_diff(xs: np.ndarray, ys: np.ndarray) -> None: check(CORPUS.boolop_numeric, xs, ys) -@given(arrays(small), arrays(small)) +@hypothesis.given(arrays(small), arrays(small)) def test_math_two_arg_diff(xs: np.ndarray, ys: np.ndarray) -> None: check(CORPUS.math_two_arg, xs, ys) -@given(arrays(positive)) +@hypothesis.given(arrays(positive)) def test_math_calls_diff(xs: np.ndarray) -> None: check(CORPUS.math_calls, xs) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_minmax_diff(xs: np.ndarray) -> None: ys = np.asarray([0.5, -1.5, 2.5] * (len(xs) // 3) + [0.0] * (len(xs) % 3)) check(CORPUS.minmax, xs, ys) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_ssa_rebind_diff(xs: np.ndarray) -> None: check(CORPUS.ssa_rebind, xs) -@given(arrays(small), arrays(small)) +@hypothesis.given(arrays(small), arrays(small)) def test_complex_expr_diff(xs: np.ndarray, ys: np.ndarray) -> None: check(CORPUS.complex_expr, xs, ys) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_casts_diff(xs: np.ndarray) -> None: check(CORPUS.casts, xs) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_compare_chain_diff(xs: np.ndarray) -> None: - vec = vectorize(CORPUS.compare_chain) + vec = array_vectorize.vectorize(CORPUS.compare_chain) got = vec(xs) expected = np.asarray([CORPUS.compare_chain(x) for x in xs]) assert np.array_equal(got, expected) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_loop_accumulate_diff(xs: np.ndarray) -> None: check(CORPUS.loop_accumulate, xs) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_loop_with_branch_diff(xs: np.ndarray) -> None: check(CORPUS.loop_with_branch, xs) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_nested_loops_diff(xs: np.ndarray) -> None: check(CORPUS.nested_loops, xs) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_helper_outer_diff(xs: np.ndarray) -> None: check(CORPUS.helper_outer, xs) -@given(arrays(small)) +@hypothesis.given(arrays(small)) def test_loop_helper_caller_diff(xs: np.ndarray) -> None: check(CORPUS.loop_helper_caller, xs) -@given(arrays(small), arrays(small)) +@hypothesis.given(arrays(small), arrays(small)) def test_second_oracle_np_vectorize(xs: np.ndarray, ys: np.ndarray) -> None: """Differential vs np.vectorize of the original.""" - vec = vectorize(CORPUS.add) + vec = array_vectorize.vectorize(CORPUS.add) oracle = np.vectorize(CORPUS.add) assert np.allclose(vec(xs, ys), oracle(xs, ys)) def test_int_dtypes_differential() -> None: - vec = vectorize(CORPUS.ssa_rebind) + vec = array_vectorize.vectorize(CORPUS.ssa_rebind) for dtype in (np.int32, np.int64): xs = np.asarray([-3, 0, 1, 7], dtype=dtype) got = vec(xs) @@ -171,16 +171,16 @@ def test_int_dtypes_differential() -> None: def test_edge_values_relu_ternary() -> None: edges = np.asarray([0.0, -0.0, 1.0, -1.0, np.inf, -np.inf, np.nan, 5e-324, -5e-324]) - vec = vectorize(CORPUS.relu) + vec = array_vectorize.vectorize(CORPUS.relu) expected = np.asarray([CORPUS.relu(float(x)) for x in edges]) assert np.allclose(vec(edges), expected, equal_nan=True) - vec2 = vectorize(CORPUS.ternary) + vec2 = array_vectorize.vectorize(CORPUS.ternary) expected2 = np.asarray([CORPUS.ternary(float(x)) for x in edges]) assert np.allclose(vec2(edges), expected2, equal_nan=True) def test_float32_differential() -> None: - vec = vectorize(CORPUS.psi) + vec = array_vectorize.vectorize(CORPUS.psi) xs = np.asarray([-1.5, -0.2, 0.0, 0.3, 2.0], dtype=np.float32) got = vec(xs) expected = np.asarray([CORPUS.psi(float(x)) for x in xs], dtype=np.float32) diff --git a/tests/features/test_fallback.py b/tests/features/test_fallback.py index 5ed5daf..3f14128 100644 --- a/tests/features/test_fallback.py +++ b/tests/features/test_fallback.py @@ -6,9 +6,9 @@ import numpy as np import pytest -from support import make_fn +import support -from array_vectorize import VectorizationError, vectorize +import array_vectorize # --------------------------------------------------------------- fallback=True @@ -20,7 +20,7 @@ def bad(x: float) -> float: return x with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(bad, fallback=True) + vec = array_vectorize.vectorize(bad, fallback=True) xs = np.asarray([5.0, -2.0]) assert np.allclose(vec(xs), [0.0, -2.0]) @@ -32,17 +32,17 @@ def bad(x: float, y: float) -> float: return x with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(bad, fallback=True) + vec = array_vectorize.vectorize(bad, fallback=True) xs = np.asarray([[3.0]]) # broadcasts against ys ys = np.asarray([[1.0, 2.0]]) assert np.allclose(vec(xs, ys), [[0.0, -1.0]]) def test_fallback_array_kwargs() -> None: - fn = make_fn(" while x < y:\n x = x + 1\n return x") + fn = support.make_fn(" while x < y:\n x = x + 1\n return x") with warnings.catch_warnings(): warnings.simplefilter("ignore") - v = vectorize(fn, fallback=True) + v = array_vectorize.vectorize(fn, fallback=True) assert np.allclose(v(x=np.asarray([1.0, 3.0])), [2.0, 3.0]) assert np.allclose(v(np.asarray([1.0, 3.0]), y=np.asarray([2.0, 4.0])), [2.0, 4.0]) @@ -56,7 +56,7 @@ def bad(x: float, scale: float = 2.0) -> float: import inspect with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(bad, fallback=True) + vec = array_vectorize.vectorize(bad, fallback=True) assert list(inspect.signature(vec).parameters) == ["x", "scale"] assert "fallback" in vec.source @@ -72,7 +72,7 @@ def bad(x: float) -> float: bad.__doc__ = 42 # type: ignore[assignment] with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(bad, fallback=True) + vec = array_vectorize.vectorize(bad, fallback=True) assert vec.__doc__ == 42 assert np.allclose(vec(np.asarray([5.0, -2.0])), [0.0, -2.0]) @@ -87,7 +87,7 @@ def bad(x: float) -> float: return x with pytest.warns(UserWarning, match="falling back"): - vec = vectorize(bad, strict=False) + vec = array_vectorize.vectorize(bad, strict=False) assert np.allclose(vec(np.asarray([5.0, -1.0])), [0.0, -1.0]) @@ -98,5 +98,5 @@ def also_bad(x: float) -> float: x = x - 1 return x - with pytest.raises(VectorizationError): - vectorize(also_bad) + with pytest.raises(array_vectorize.VectorizationError): + array_vectorize.vectorize(also_bad) diff --git a/tests/features/test_protect_domains.py b/tests/features/test_protect_domains.py index c2f1306..7577d0d 100644 --- a/tests/features/test_protect_domains.py +++ b/tests/features/test_protect_domains.py @@ -5,16 +5,16 @@ import math import numpy as np -from support import make_fn +import support -from array_vectorize import vectorize +import array_vectorize # ---------------------------------------------------------- dead-lane clamping def test_protect_domains_clamps_dead_lanes() -> None: - fn = make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") - vec = vectorize(fn, protect_domains=True) + fn = support.make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") + vec = array_vectorize.vectorize(fn, protect_domains=True) assert "xp.logical_not(x > 0)" in vec.source # dead-lane guard xs = np.asarray([-1.0, 0.5, 4.0]) with np.errstate(invalid="ignore", divide="ignore"): @@ -23,8 +23,8 @@ def test_protect_domains_clamps_dead_lanes() -> None: def test_protect_domains_no_warning() -> None: - fn = make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") - vec = vectorize(fn, protect_domains=True) + fn = support.make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") + vec = array_vectorize.vectorize(fn, protect_domains=True) xs = np.asarray([-100.0, -1.0]) with np.errstate(all="raise"): # any warning becomes an error got = vec(xs) @@ -32,22 +32,26 @@ def test_protect_domains_no_warning() -> None: def test_protect_domains_no_warning_on_poles() -> None: - vec = vectorize(make_fn(" return math.log(x) if x > 0 else 0.0"), protect_domains=True) + vec = array_vectorize.vectorize( + support.make_fn(" return math.log(x) if x > 0 else 0.0"), protect_domains=True + ) with np.errstate(all="raise"): got = vec(np.asarray([-1.0, 1.0])) assert np.allclose(got, [0.0, 0.0]) def test_protect_domains_open_boundaries_never_warn() -> None: - vec = vectorize(make_fn(" return math.log1p(x) if x > -1 else 0.0"), protect_domains=True) + vec = array_vectorize.vectorize( + support.make_fn(" return math.log1p(x) if x > -1 else 0.0"), protect_domains=True + ) with np.errstate(all="raise"): got = vec(np.asarray([-1.0, 0.0])) assert np.allclose(got, [0.0, 0.0]) def test_protect_domains_two_sided_clamp() -> None: - fn = make_fn(" if x != 0:\n return math.atanh(x)\n return 0.0") - vec = vectorize(fn, protect_domains=True) + fn = support.make_fn(" if x != 0:\n return math.atanh(x)\n return 0.0") + vec = array_vectorize.vectorize(fn, protect_domains=True) # dead-lane where-clamps (preserve -0.0 and live values, unlike maximum) assert "xp.logical_not" in vec.source assert "xp.maximum" not in vec.source @@ -56,8 +60,8 @@ def test_protect_domains_two_sided_clamp() -> None: def test_protect_domains_atanh_dead_lanes() -> None: - vec = vectorize( - make_fn(" return math.atanh(x) if -1 < x < 1 else 0.0"), protect_domains=True + vec = array_vectorize.vectorize( + support.make_fn(" return math.atanh(x) if -1 < x < 1 else 0.0"), protect_domains=True ) with np.errstate(all="raise"): got = vec(np.asarray([2.0, 0.5, -2.0])) @@ -68,23 +72,25 @@ def test_protect_domains_atanh_dead_lanes() -> None: def test_protect_domains_never_changes_results() -> None: - fn = make_fn(" if -1 <= x <= 1:\n return math.asin(x)\n return 0.0") - vec = vectorize(fn, protect_domains=True) + fn = support.make_fn(" if -1 <= x <= 1:\n return math.asin(x)\n return 0.0") + vec = array_vectorize.vectorize(fn, protect_domains=True) xs = np.asarray([-2.0, -0.5, 0.0, 0.5, 2.0]) expected = np.asarray([0.0, math.asin(-0.5), 0.0, math.asin(0.5), 0.0]) assert np.allclose(vec(xs), expected) def test_protect_domains_preserves_signed_zero() -> None: - vec = vectorize( - make_fn(" return math.copysign(1.0, math.sqrt(x)) if x <= 0 else 5.0"), + vec = array_vectorize.vectorize( + support.make_fn(" return math.copysign(1.0, math.sqrt(x)) if x <= 0 else 5.0"), protect_domains=True, ) assert vec(np.asarray([-0.0]))[0] == -1.0 def test_protect_domains_preserves_subnormal_live_lanes() -> None: - vec = vectorize(make_fn(" return math.log(x) if x > 0 else 0.0"), protect_domains=True) + vec = array_vectorize.vectorize( + support.make_fn(" return math.log(x) if x > 0 else 0.0"), protect_domains=True + ) got = vec(np.asarray([1e-310])) assert got[0] == math.log(1e-310) @@ -93,22 +99,22 @@ def test_protect_domains_preserves_subnormal_live_lanes() -> None: def test_protect_domains_off_by_default() -> None: - fn = make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") - vec = vectorize(fn) + fn = support.make_fn(" if x > 0:\n return math.sqrt(x)\n return 0.0") + vec = array_vectorize.vectorize(fn) assert "xp.maximum(x, 0.0)" not in vec.source def test_protect_domains_top_level_calls_not_clamped() -> None: - fn = make_fn(" return math.sqrt(abs(x))") - vec = vectorize(fn, protect_domains=True) + fn = support.make_fn(" return math.sqrt(abs(x))") + vec = array_vectorize.vectorize(fn, protect_domains=True) assert "maximum" not in vec.source def test_protect_domains_conditions_not_clamped() -> None: # a partial call inside a condition must not be clamped (results # never change); isnan(sqrt(x)) is a condition here - fn = make_fn(" return 1.0 if math.isnan(math.sqrt(x)) else 0.0") - vec = vectorize(fn, protect_domains=True) + fn = support.make_fn(" return 1.0 if math.isnan(math.sqrt(x)) else 0.0") + vec = array_vectorize.vectorize(fn, protect_domains=True) with np.errstate(invalid="ignore"): got = vec(np.asarray([-1.0, 4.0])) assert np.array_equal(got, [1.0, 0.0]) @@ -118,8 +124,8 @@ def test_protect_domains_conditions_not_clamped() -> None: def test_protect_domains_follows_branch_bindings() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " if x > 0:\n y = math.sqrt(x)\n else:\n y = 0.0\n return y" ), protect_domains=True, @@ -133,8 +139,8 @@ def test_protect_domains_multiple_uses_or_liveness() -> None: # y_1 is used in two wheres with different (overlapping) conditions: # only lanes dead for BOTH uses are clampable; live lanes keep exact # values. Live region = (x > 0) or (x >= 0) = x >= 0; dead = x < 0. - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " y = math.sqrt(x)\n" " a = y if x > 0 else 0.0\n" " b = y if x >= 0 else 1.0\n" @@ -148,8 +154,8 @@ def test_protect_domains_multiple_uses_or_liveness() -> None: def test_protect_domains_self_reference_guard() -> None: - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " if x < 0:\n" " return 0.0\n" " y = math.sqrt(x)\n" @@ -165,8 +171,8 @@ def test_protect_domains_self_reference_guard() -> None: def test_protect_domains_forward_reference_guard() -> None: # the liveness condition references z, which is bound after the sqrt: # the clamp must be skipped (not crash), values stay exact - vec = vectorize( - make_fn( + vec = array_vectorize.vectorize( + support.make_fn( " y = math.sqrt(x)\n" " if x > 0:\n" " z = x\n" @@ -188,8 +194,10 @@ def test_protect_domains_loop_body_binding_guard() -> None: " c = x > 0\n" " return y if c else 0.0" ) - plain = vectorize(make_fn(body))(np.asarray([4.0])) - protected = vectorize(make_fn(body), protect_domains=True)(np.asarray([4.0])) + plain = array_vectorize.vectorize(support.make_fn(body))(np.asarray([4.0])) + protected = array_vectorize.vectorize(support.make_fn(body), protect_domains=True)( + np.asarray([4.0]) + ) assert plain[0] == 2.0 assert protected[0] == 2.0 @@ -198,11 +206,11 @@ def test_protect_domains_propagates_to_helpers() -> None: # the helper's own internal where gets protected; note that a helper # called only inside a caller's where-branch cannot know its lanes are # dead (cross-function liveness is a documented limitation) - fn = make_fn( + fn = support.make_fn( " return helper(x) * 2.0", extra="def helper(v):\n return math.sqrt(v) if v > 0 else 0.0\n", ) - vec = vectorize(fn, protect_domains=True) + vec = array_vectorize.vectorize(fn, protect_domains=True) with np.errstate(all="raise"): got = vec(np.asarray([-1.0, 4.0])) assert np.allclose(got, [0.0, 4.0]) diff --git a/tests/features/test_verify.py b/tests/features/test_verify.py index 77dac32..9a18ea3 100644 --- a/tests/features/test_verify.py +++ b/tests/features/test_verify.py @@ -4,25 +4,25 @@ import numpy as np import pytest -from support import make_fn +import support -from array_vectorize import VectorizationError, vectorize -from array_vectorize.verify import verify_match +import array_vectorize +from array_vectorize import verify # ------------------------------------------------------- vectorize(verify=...) def test_verify_passes_on_match() -> None: - fn = make_fn(" if x < 0:\n return 0.0\n return x * math.exp(-x)") + fn = support.make_fn(" if x < 0:\n return 0.0\n return x * math.exp(-x)") xs = np.asarray([-1.0, 0.0, 2.0]) - vec = vectorize(fn, verify=(xs,)) + vec = array_vectorize.vectorize(fn, verify=(xs,)) assert np.allclose(vec(xs), np.where(xs < 0, 0.0, xs * np.exp(-xs))) def test_verify_nan_aware() -> None: - fn = make_fn(" return math.sqrt(x)") + fn = support.make_fn(" return math.sqrt(x)") xs = np.asarray([-1.0, 4.0]) - vec = vectorize(fn, verify=(xs,)) # NaN lane matches NaN lane + vec = array_vectorize.vectorize(fn, verify=(xs,)) # NaN lane matches NaN lane with np.errstate(invalid="ignore"): got = vec(xs) assert np.isnan(got[0]) @@ -34,36 +34,38 @@ def test_verify_detects_mismatch() -> None: # a function whose vectorization is correct still passes; instead check # the failure path with a fallback-wrapped lie is out of scope. Use a # shape mismatch instead. - fn = make_fn(" return x + 0.0") + fn = support.make_fn(" return x + 0.0") xs = np.asarray([[1.0, 2.0], [3.0, 4.0]]) # 2-D is fine... - vec = vectorize(fn, verify=(xs,)) + vec = array_vectorize.vectorize(fn, verify=(xs,)) assert vec(xs).shape == (2, 2) def test_verify_passes_scalar_arguments_through() -> None: - fn = make_fn(" return x + y") - vec = vectorize(fn, verify=(np.asarray([1.0]), 3)) + fn = support.make_fn(" return x + y") + vec = array_vectorize.vectorize(fn, verify=(np.asarray([1.0]), 3)) assert vec(np.asarray([1.0]), 3)[0] == 4.0 def test_verify_constant_return() -> None: - vec = vectorize(make_fn(" return 3.0"), verify=(np.asarray([1.0, 2.0]),)) + vec = array_vectorize.vectorize( + support.make_fn(" return 3.0"), verify=(np.asarray([1.0, 2.0]),) + ) got = vec(np.asarray([1.0, 2.0])) assert np.allclose(np.broadcast_to(np.asarray(got), (2,)), [3.0, 3.0]) def test_verify_requires_array() -> None: - fn = make_fn(" return x + 1.0") - with pytest.raises(VectorizationError, match="at least one Array API array"): - vectorize(fn, verify=(1.0,)) + fn = support.make_fn(" return x + 1.0") + with pytest.raises(array_vectorize.VectorizationError, match="at least one Array API array"): + array_vectorize.vectorize(fn, verify=(1.0,)) # ---------------------------------------------------------------- verify_match def test_verify_detects_bool_miscompile() -> None: - with pytest.raises(VectorizationError, match="verification failed"): - verify_match( + with pytest.raises(array_vectorize.VectorizationError, match="verification failed"): + verify.verify_match( lambda x: np.array([False, True]), # wrong: should be [0, 2] lambda x: 2 if x > 0 else 0, (np.array([-1.0, 1.0]),), @@ -72,14 +74,14 @@ def test_verify_detects_bool_miscompile() -> None: def test_verify_detects_wrong_constant_output() -> None: # a "vectorized" function returning the wrong constant must be caught - with pytest.raises(VectorizationError, match="verification failed"): - verify_match(lambda x: np.array([1.0]), lambda x: x * 5.0, (np.array([2.0]),)) + with pytest.raises(array_vectorize.VectorizationError, match="verification failed"): + verify.verify_match(lambda x: np.array([1.0]), lambda x: x * 5.0, (np.array([2.0]),)) def test_verify_accepts_matching_infinities() -> None: - verify_match(lambda x: x, lambda x: x, (np.array([np.inf, -np.inf]),)) + verify.verify_match(lambda x: x, lambda x: x, (np.array([np.inf, -np.inf]),)) def test_verify_rejects_inf_vs_finite() -> None: - with pytest.raises(VectorizationError, match="verification failed"): - verify_match(lambda x: np.array([0.0]), lambda x: float("inf"), (np.array([1.0]),)) + with pytest.raises(array_vectorize.VectorizationError, match="verification failed"): + verify.verify_match(lambda x: np.array([0.0]), lambda x: float("inf"), (np.array([1.0]),)) diff --git a/tests/golden/test_golden.py b/tests/golden/test_golden.py index 9bd9528..b81dbaf 100644 --- a/tests/golden/test_golden.py +++ b/tests/golden/test_golden.py @@ -3,16 +3,15 @@ from __future__ import annotations import ast -from pathlib import Path +import pathlib import corpus import numpy as np import pytest -from corpus import GOLDEN_NAMES -from array_vectorize import vectorize +import array_vectorize -GOLDEN_DIR = Path(__file__).parent / "cases" +GOLDEN_DIR = pathlib.Path(__file__).parent / "cases" #: pinned-variant goldens (vectorize(namespace=np)): the namespace object #: never appears in the source text, only the ``xp = _namespace`` binding @@ -21,19 +20,19 @@ def _generated_source(name: str) -> str: fn = getattr(corpus, name) - return vectorize(fn).source + return array_vectorize.vectorize(fn).source def _generated_pinned_source(name: str) -> str: fn = getattr(corpus, name) - return vectorize(fn, namespace=np).source + return array_vectorize.vectorize(fn, namespace=np).source -def _golden_path(name: str) -> Path: +def _golden_path(name: str) -> pathlib.Path: return GOLDEN_DIR / f"{name}.py" -def _pinned_golden_path(name: str) -> Path: +def _pinned_golden_path(name: str) -> pathlib.Path: return GOLDEN_DIR / f"{name}_pinned.py" @@ -41,7 +40,7 @@ def _tree(source: str) -> str: return ast.dump(ast.parse(source)) -@pytest.mark.parametrize("name", GOLDEN_NAMES) +@pytest.mark.parametrize("name", corpus.GOLDEN_NAMES) def test_golden_source(name: str, request: pytest.FixtureRequest) -> None: source = _generated_source(name) path = _golden_path(name) diff --git a/tests/perf/test_bench.py b/tests/perf/test_bench.py index 58e04b5..ff829a4 100644 --- a/tests/perf/test_bench.py +++ b/tests/perf/test_bench.py @@ -23,9 +23,8 @@ import numpy as np import pytest -from array_vectorize import vectorize -from array_vectorize.runtime.arith import _ARITH_OPS, _vec_arith -from array_vectorize.runtime.minmax import _vec_minmax +import array_vectorize +from array_vectorize.runtime import arith, minmax pytestmark = pytest.mark.slow @@ -40,7 +39,7 @@ ] SIZES = [1_000, 100_000, 1_000_000] -_add_id = _ARITH_OPS.index("add") +_add_id = arith._ARITH_OPS.index("add") @pytest.mark.parametrize(("name", "kernel", "nargs"), KERNELS, ids=[k[0] for k in KERNELS]) @@ -52,7 +51,7 @@ def test_kernel_vs_oracle(benchmark, name, kernel, nargs, size) -> None: y = rng.standard_normal(size) args = (x, y) if nargs == 2 else (x,) small = tuple(a[:100] for a in args) - vec = vectorize(kernel) + vec = array_vectorize.vectorize(kernel) oracle = np.vectorize(kernel) assert np.allclose(vec(*small), oracle(*small), equal_nan=True) @@ -91,7 +90,7 @@ def subject(x, y=2.0): acc = acc + math.sqrt(x) return acc - vec = benchmark(vectorize, subject) + vec = benchmark(array_vectorize.vectorize, subject) assert vec(np.asarray([4.0]))[0] == pytest.approx(6.0) @@ -103,7 +102,7 @@ def test_helper_dispatch_overhead(benchmark) -> None: def run() -> None: for _ in range(10): - _vec_arith(xp, _add_id, a, b) + arith._vec_arith(xp, _add_id, a, b) benchmark(run) if not benchmark.disabled: @@ -122,7 +121,7 @@ def test_minmax_dispatch_overhead(benchmark) -> None: def run() -> None: for _ in range(10): - _vec_minmax(xp, False, a, b) + minmax._vec_minmax(xp, False, a, b) benchmark(run) if not benchmark.disabled: diff --git a/tests/perf/test_speed_gate.py b/tests/perf/test_speed_gate.py index 8a4afc3..133ef8a 100644 --- a/tests/perf/test_speed_gate.py +++ b/tests/perf/test_speed_gate.py @@ -21,7 +21,7 @@ import numpy as np import pytest -from array_vectorize import vectorize +import array_vectorize pytestmark = pytest.mark.slow @@ -48,7 +48,7 @@ def test_faster_than_np_vectorize() -> None: # a single-pass function maximizes the ratio: the comparison measures # the Python-call overhead that vectorization eliminates (multi-op or # exp-heavy bodies are Amdahl-limited by the array work itself). - vec = vectorize(lambda v: v + 1.0) + vec = array_vectorize.vectorize(lambda v: v + 1.0) oracle = np.vectorize(lambda v: v + 1.0) expected = oracle(x[:1000]) # warm up + sanity diff --git a/tests/support.py b/tests/support.py index c312a01..0156d64 100644 --- a/tests/support.py +++ b/tests/support.py @@ -12,17 +12,17 @@ import importlib.util import itertools +import pathlib import tempfile from collections.abc import Callable -from pathlib import Path from typing import Any import corpus -from array_vectorize import vectorize +import array_vectorize _tmp = tempfile.TemporaryDirectory(prefix="vec_support_") -TMPDIR = Path(_tmp.name) +TMPDIR = pathlib.Path(_tmp.name) _seq = itertools.count() @@ -64,13 +64,13 @@ def make_module(source: str) -> Any: def vec_of(body: str, extra: str = "") -> tuple[Callable[..., Any], Callable[..., Any]]: fn = make_fn(body, extra) - return vectorize(fn), fn + return array_vectorize.vectorize(fn), fn def just_vec(body: str, extra: str = "") -> Callable[..., Any]: - return vectorize(make_fn(body, extra)) + return array_vectorize.vectorize(make_fn(body, extra)) def vfn(name: str) -> Callable[..., Any]: """Vectorize the corpus function ``name`` (tests/corpus.py).""" - return vectorize(getattr(corpus, name)) + return array_vectorize.vectorize(getattr(corpus, name)) diff --git a/tests/test_agent_layout.py b/tests/test_agent_layout.py index 6eae51e..6d5ea98 100644 --- a/tests/test_agent_layout.py +++ b/tests/test_agent_layout.py @@ -6,12 +6,12 @@ from __future__ import annotations +import pathlib import re -from pathlib import Path import pytest -ROOT = Path(__file__).parents[1] +ROOT = pathlib.Path(__file__).parents[1] SKILLS = ROOT / ".agents" / "skills" CLAUDE_SKILLS = ROOT / ".claude" / "skills" CLAUDE_AGENTS = ROOT / ".claude" / "agents" @@ -22,7 +22,7 @@ SKILL_NAME = re.compile(r"[a-z0-9]+(-[a-z0-9]+)*") -def _frontmatter(path: Path) -> dict[str, str | list[str]]: +def _frontmatter(path: pathlib.Path) -> dict[str, str | list[str]]: """Top-level keys of a markdown file's YAML frontmatter (scalars, lists).""" lines = path.read_text().splitlines() assert lines[0] == "---", f"{path} has no frontmatter" @@ -47,7 +47,7 @@ def test_there_are_skills() -> None: @pytest.mark.parametrize("skill", SKILL_DIRS, ids=lambda p: p.name) -def test_skill_uses_standard_frontmatter(skill: Path) -> None: +def test_skill_uses_standard_frontmatter(skill: pathlib.Path) -> None: fields = _frontmatter(skill / "SKILL.md") assert set(fields) <= STANDARD_FIELDS, set(fields) - STANDARD_FIELDS assert fields["name"] == skill.name @@ -66,7 +66,7 @@ def test_claude_skills_are_symlinks_to_agents_skills() -> None: @pytest.mark.parametrize("agent", sorted(CLAUDE_AGENTS.glob("*.md")), ids=lambda p: p.name) -def test_claude_subagents_preload_shared_skills(agent: Path) -> None: +def test_claude_subagents_preload_shared_skills(agent: pathlib.Path) -> None: skills = _frontmatter(agent).get("skills", []) assert skills, f"{agent} should wrap a skill from .agents/skills" assert isinstance(skills, list) diff --git a/tests/unit/test_builtins.py b/tests/unit/test_builtins.py index b5fceb3..8eb5910 100644 --- a/tests/unit/test_builtins.py +++ b/tests/unit/test_builtins.py @@ -6,34 +6,27 @@ import array_api_strict as xps -from array_vectorize.frontend.tables import ( - BUILTIN_CASTS, - BUILTIN_FOLDS, - BUILTIN_UNARY, - MATH_CONSTS, - MATH_FUNCS, - MATH_SPECIAL, -) +from array_vectorize.frontend import tables def test_every_math_func_exists_on_strict() -> None: - for xp_name, _, _ in MATH_FUNCS.values(): + for xp_name, _, _ in tables.MATH_FUNCS.values(): assert hasattr(xps, xp_name), f"xp.{xp_name} missing on array-api-strict" def test_every_builtin_target_exists_on_strict() -> None: - for xp_name in BUILTIN_UNARY.values(): + for xp_name in tables.BUILTIN_UNARY.values(): assert hasattr(xps, xp_name) - for xp_name in BUILTIN_FOLDS.values(): + for xp_name in tables.BUILTIN_FOLDS.values(): assert hasattr(xps, xp_name) - for dtype in BUILTIN_CASTS.values(): + for dtype in tables.BUILTIN_CASTS.values(): assert hasattr(xps, dtype), f"xp.{dtype} missing on array-api-strict" - for dtype in MATH_SPECIAL.values(): + for dtype in tables.MATH_SPECIAL.values(): assert hasattr(xps, dtype) def test_math_consts_match_math_module() -> None: - for name, value in MATH_CONSTS.items(): + for name, value in tables.MATH_CONSTS.items(): if name == "nan": assert math.isnan(value) else: diff --git a/tests/unit/test_codegen.py b/tests/unit/test_codegen.py index 5fa638b..0a5db38 100644 --- a/tests/unit/test_codegen.py +++ b/tests/unit/test_codegen.py @@ -4,28 +4,15 @@ import ast -from array_vectorize.codegen import generate_source -from array_vectorize.frontend.info import Param -from array_vectorize.ir import ( - Binding, - BinOp, - Call, - Compare, - DType, - Literal, - Logical, - Program, - Ref, - UnaryOp, - Where, -) -from array_vectorize.lower.types import LoweredFunction +from array_vectorize import codegen, ir +from array_vectorize.frontend import info +from array_vectorize.lower import types def make_lowered( - params: list[Param], program: Program, name: str = "f", docstring: str | None = None -) -> LoweredFunction: - return LoweredFunction( + params: list[info.Param], program: ir.Program, name: str = "f", docstring: str | None = None +) -> types.LoweredFunction: + return types.LoweredFunction( program=program, name=name, params=params, @@ -40,18 +27,18 @@ def make_lowered( ) -def gen(program: Program, params: list[Param] | None = None) -> str: - params = params if params is not None else [Param("x", "arg")] - return generate_source(make_lowered(params, program), program) +def gen(program: ir.Program, params: list[info.Param] | None = None) -> str: + params = params if params is not None else [info.Param("x", "arg")] + return codegen.generate_source(make_lowered(params, program), program) def test_simple_return() -> None: - src = gen(Program(("x",), (), BinOp("add", Ref("x"), Literal(1, "int")))) + src = gen(ir.Program(("x",), (), ir.BinOp("add", ir.Ref("x"), ir.Literal(1, "int")))) assert "return x + 1" in src def test_module_shape() -> None: - src = gen(Program(("x",), (), Ref("x"))) + src = gen(ir.Program(("x",), (), ir.Ref("x"))) assert src.startswith("from array_api_compat import array_namespace") assert "def f_vec(x):" in src assert '"""\n (array-vectorized) no docstring on the scalar original.' in src @@ -62,47 +49,47 @@ def test_module_shape() -> None: def test_module_shape_with_docstring() -> None: - program = Program(("x",), (), Ref("x")) - lowered = make_lowered([Param("x", "arg")], program, docstring="Double x.") - src = generate_source(lowered, program) + program = ir.Program(("x",), (), ir.Ref("x")) + lowered = make_lowered([info.Param("x", "arg")], program, docstring="Double x.") + src = codegen.generate_source(lowered, program) assert "(array-vectorized) Double x." in src assert "no docstring on the scalar original" not in src def test_operators_stay_operators() -> None: src = gen( - Program( + ir.Program( ("x", "y"), (), - BinOp( + ir.BinOp( "sub", - BinOp("mul", Ref("x"), Ref("y")), - BinOp("div", Ref("x"), Ref("y")), + ir.BinOp("mul", ir.Ref("x"), ir.Ref("y")), + ir.BinOp("div", ir.Ref("x"), ir.Ref("y")), ), ), - [Param("x", "arg"), Param("y", "arg")], + [info.Param("x", "arg"), info.Param("y", "arg")], ) assert "return x * y - x / y" in src def test_div_pow_mod_floordiv_become_xp_calls() -> None: - src = gen(Program(("x",), (), BinOp("pow", Ref("x"), Literal(2, "int")))) + src = gen(ir.Program(("x",), (), ir.BinOp("pow", ir.Ref("x"), ir.Literal(2, "int")))) assert "xp.pow(x, 2)" in src - src = gen(Program(("x",), (), BinOp("floordiv", Ref("x"), Literal(2, "int")))) + src = gen(ir.Program(("x",), (), ir.BinOp("floordiv", ir.Ref("x"), ir.Literal(2, "int")))) assert "xp.floor_divide(x, 2)" in src - src = gen(Program(("x",), (), BinOp("mod", Ref("x"), Literal(2, "int")))) + src = gen(ir.Program(("x",), (), ir.BinOp("mod", ir.Ref("x"), ir.Literal(2, "int")))) assert "xp.remainder(x, 2)" in src def test_logical_not_where() -> None: src = gen( - Program( + ir.Program( ("x",), (), - Where( - UnaryOp("not", Compare("gt", Ref("x"), Literal(0, "int"))), - Literal(1.0, "float"), - Ref("x"), + ir.Where( + ir.UnaryOp("not", ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int"))), + ir.Literal(1.0, "float"), + ir.Ref("x"), ), ) ) @@ -111,15 +98,15 @@ def test_logical_not_where() -> None: def test_logical_and_folded_pairwise() -> None: src = gen( - Program( + ir.Program( ("x",), (), - Logical( + ir.Logical( "and", ( - Compare("lt", Literal(0, "int"), Ref("x")), - Compare("lt", Ref("x"), Literal(1, "int")), - Compare("lt", Ref("x"), Literal(2, "int")), + ir.Compare("lt", ir.Literal(0, "int"), ir.Ref("x")), + ir.Compare("lt", ir.Ref("x"), ir.Literal(1, "int")), + ir.Compare("lt", ir.Ref("x"), ir.Literal(2, "int")), ), ), ) @@ -128,27 +115,33 @@ def test_logical_and_folded_pairwise() -> None: def test_inf_nan_literals() -> None: - src = gen(Program(("x",), (), BinOp("add", Ref("x"), Literal(float("inf"), "float")))) + src = gen( + ir.Program(("x",), (), ir.BinOp("add", ir.Ref("x"), ir.Literal(float("inf"), "float"))) + ) assert "x + xp.inf" in src - src = gen(Program(("x",), (), BinOp("add", Ref("x"), Literal(float("-inf"), "float")))) + src = gen( + ir.Program(("x",), (), ir.BinOp("add", ir.Ref("x"), ir.Literal(float("-inf"), "float"))) + ) assert "x + -xp.inf" in src - src = gen(Program(("x",), (), BinOp("add", Ref("x"), Literal(float("nan"), "float")))) + src = gen( + ir.Program(("x",), (), ir.BinOp("add", ir.Ref("x"), ir.Literal(float("nan"), "float"))) + ) assert "x + xp.nan" in src def test_casts() -> None: - src = gen(Program(("x",), (), Call("astype", (Ref("x"), DType("int64"))))) + src = gen(ir.Program(("x",), (), ir.Call("astype", (ir.Ref("x"), ir.DType("int64"))))) assert "xp.astype(x, xp.int64)" in src def test_bindings_and_temp() -> None: - program = Program( + program = ir.Program( ("x",), ( - Binding("t_1", Call("sqrt", (Ref("x"),))), - Binding("a", BinOp("add", Ref("t_1"), Ref("t_1"))), + ir.Binding("t_1", ir.Call("sqrt", (ir.Ref("x"),))), + ir.Binding("a", ir.BinOp("add", ir.Ref("t_1"), ir.Ref("t_1"))), ), - Ref("a"), + ir.Ref("a"), ) src = gen(program) assert "t_1 = xp.sqrt(x)" in src @@ -160,20 +153,20 @@ def test_bindings_and_temp() -> None: def test_defaults_and_kwonly() -> None: params = [ - Param("x", "arg"), - Param("scale", "arg", 2.0, True), - Param("flag", "kwonly", True, True), + info.Param("x", "arg"), + info.Param("scale", "arg", 2.0, True), + info.Param("flag", "kwonly", True, True), ] - program = Program(("x", "scale", "flag"), (), Ref("x")) + program = ir.Program(("x", "scale", "flag"), (), ir.Ref("x")) src = gen(program, params) assert "def f_vec(x, scale=2.0, *, flag=True):" in src def test_hidden_params_kwonly_none() -> None: - lowered = LoweredFunction( - program=Program(("x", "ARR"), (), Ref("ARR")), + lowered = types.LoweredFunction( + program=ir.Program(("x", "ARR"), (), ir.Ref("ARR")), name="f", - params=[Param("x", "arg")], + params=[info.Param("x", "arg")], param_names=["x"], hidden_params=[("ARR", "placeholder")], helpers=[], @@ -183,22 +176,22 @@ def test_hidden_params_kwonly_none() -> None: source="def f(x):\n return ARR", docstring=None, ) - src = generate_source(lowered, lowered.program) + src = codegen.generate_source(lowered, lowered.program) assert "def f_vec(x, *, ARR=None):" in src assert "xp = array_namespace(x, ARR)" in src def test_posonly_marker() -> None: - params = [Param("x", "posonly"), Param("y", "arg")] - program = Program(("x", "y"), (), Ref("x")) + params = [info.Param("x", "posonly"), info.Param("y", "arg")] + program = ir.Program(("x", "y"), (), ir.Ref("x")) src = gen(program, params) assert "def f_vec(x, /, y):" in src def test_reserved_param_kept_namespace_renamed() -> None: - params = [Param("xp", "arg")] - lowered = LoweredFunction( - program=Program(("xp",), (), Ref("xp")), + params = [info.Param("xp", "arg")] + lowered = types.LoweredFunction( + program=ir.Program(("xp",), (), ir.Ref("xp")), name="f", params=params, param_names=["xp"], @@ -210,28 +203,32 @@ def test_reserved_param_kept_namespace_renamed() -> None: source="def f(xp):\n return xp", docstring=None, ) - src = generate_source(lowered, lowered.program) + src = codegen.generate_source(lowered, lowered.program) assert "def f_vec(xp):" in src assert "xp_1 = array_namespace(" in src def test_generated_source_is_valid_python() -> None: - program = Program( + program = ir.Program( ("x", "y"), ( - Binding("a", Call("sqrt", (BinOp("add", Ref("x"), Ref("y")),))), - Binding( + ir.Binding("a", ir.Call("sqrt", (ir.BinOp("add", ir.Ref("x"), ir.Ref("y")),))), + ir.Binding( "b", - Where(Compare("gt", Ref("a"), Literal(0, "int")), Ref("a"), Literal(0.0, "float")), + ir.Where( + ir.Compare("gt", ir.Ref("a"), ir.Literal(0, "int")), + ir.Ref("a"), + ir.Literal(0.0, "float"), + ), ), ), - BinOp("add", Ref("b"), Literal(1, "int")), + ir.BinOp("add", ir.Ref("b"), ir.Literal(1, "int")), ) - src = gen(program, [Param("x", "arg"), Param("y", "arg")]) + src = gen(program, [info.Param("x", "arg"), info.Param("y", "arg")]) ast.parse(src) # must not raise def test_no_numpy_leakage() -> None: - src = gen(Program(("x",), (), Call("sqrt", (Ref("x"),)))) + src = gen(ir.Program(("x",), (), ir.Call("sqrt", (ir.Ref("x"),)))) assert "np." not in src assert "numpy" not in src diff --git a/tests/unit/test_docstring.py b/tests/unit/test_docstring.py index 3c27d70..b7b4c6a 100644 --- a/tests/unit/test_docstring.py +++ b/tests/unit/test_docstring.py @@ -2,20 +2,11 @@ from __future__ import annotations -from inspect import cleandoc +import inspect import pytest -from array_vectorize.codegen.docstring import ( - DOCSTRING_PREFIX, - LEAD_IN, - NO_DOCSTRING_SUMMARY, - build_docstring, - extract_scalar_source, - is_numpy_style, - prefixed_summary, -) -from array_vectorize.codegen.emitter import _docstring_value +from array_vectorize.codegen import docstring, emitter SOURCE = ( "def clamp(x, lo=0.0, hi=1.0):\n" @@ -41,36 +32,36 @@ def test_prefix_on_summary_line() -> None: - doc = build_docstring(GOOGLE_DOC, SOURCE) - assert doc.splitlines()[0] == f"{DOCSTRING_PREFIX}Clamp ``x``." + doc = docstring.build_docstring(GOOGLE_DOC, SOURCE) + assert doc.splitlines()[0] == f"{docstring.DOCSTRING_PREFIX}Clamp ``x``." assert "Args:" in doc def test_no_docstring_synthesizes_summary() -> None: - doc = build_docstring(None, SOURCE) - assert doc.splitlines()[0] == DOCSTRING_PREFIX + NO_DOCSTRING_SUMMARY + doc = docstring.build_docstring(None, SOURCE) + assert doc.splitlines()[0] == docstring.DOCSTRING_PREFIX + docstring.NO_DOCSTRING_SUMMARY @pytest.mark.parametrize("blank", ["", " ", "\n\n"]) def test_blank_docstring_treated_as_absent(blank: str) -> None: - assert build_docstring(blank, SOURCE) == build_docstring(None, SOURCE) + assert docstring.build_docstring(blank, SOURCE) == docstring.build_docstring(None, SOURCE) def test_google_notes_section_shape() -> None: - doc = build_docstring(GOOGLE_DOC, SOURCE) + doc = docstring.build_docstring(GOOGLE_DOC, SOURCE) lines = doc.splitlines() assert "Notes:" in lines lead = lines[lines.index("Notes:") + 1] - assert lead == f" {LEAD_IN}" + assert lead == f" {docstring.LEAD_IN}" assert lines[lines.index(lead) + 2] == " " * 8 + "def clamp(x, lo=0.0, hi=1.0):" def test_numpy_notes_section_shape() -> None: - doc = build_docstring(NUMPY_DOC, SOURCE) + doc = docstring.build_docstring(NUMPY_DOC, SOURCE) lines = doc.splitlines() assert "Notes" in lines assert lines[lines.index("Notes") + 1] == "-----" - assert lines[lines.index("Notes") + 2] == LEAD_IN + assert lines[lines.index("Notes") + 2] == docstring.LEAD_IN assert lines[lines.index("Notes") + 4] == " " * 4 + "def clamp(x, lo=0.0, hi=1.0):" @@ -98,51 +89,51 @@ def test_numpy_notes_section_shape() -> None: ], ) def test_round_trip(doc: str) -> None: - assert extract_scalar_source(build_docstring(doc, SOURCE)) == SOURCE + assert docstring.extract_scalar_source(docstring.build_docstring(doc, SOURCE)) == SOURCE def test_google_notes_merged_not_duplicated() -> None: - doc = build_docstring(GOOGLE_NOTES_DOC, SOURCE) + doc = docstring.build_docstring(GOOGLE_NOTES_DOC, SOURCE) assert [line for line in doc.splitlines() if line.strip() == "Notes:"] == ["Notes:"] # the pre-existing note stays, before our marker - assert doc.index("Pre-existing note.") < doc.index(LEAD_IN) + assert doc.index("Pre-existing note.") < doc.index(docstring.LEAD_IN) def test_google_notes_in_middle_inserts_before_next_header() -> None: - doc = build_docstring(GOOGLE_NOTES_MIDDLE_DOC, SOURCE) + doc = docstring.build_docstring(GOOGLE_NOTES_MIDDLE_DOC, SOURCE) # our source block lands inside the Notes section, before Returns: assert doc.index("def clamp(x") < doc.index("Returns:") - assert doc.index(LEAD_IN) < doc.index("Returns:") + assert doc.index(docstring.LEAD_IN) < doc.index("Returns:") def test_numpy_notes_merged_not_duplicated() -> None: - doc = build_docstring(NUMPY_NOTES_DOC, SOURCE) + doc = docstring.build_docstring(NUMPY_NOTES_DOC, SOURCE) titles = [i for i, line in enumerate(doc.splitlines()) if line == "Notes"] assert len(titles) == 1 - assert doc.index("Pre-existing note.") < doc.index(LEAD_IN) + assert doc.index("Pre-existing note.") < doc.index(docstring.LEAD_IN) def test_numpy_notes_in_middle_inserts_before_next_header() -> None: - doc = build_docstring(NUMPY_NOTES_MIDDLE_DOC, SOURCE) + doc = docstring.build_docstring(NUMPY_NOTES_MIDDLE_DOC, SOURCE) assert doc.index("def clamp(x") < doc.index("Examples") - assert doc.index(LEAD_IN) < doc.index("Examples") + assert doc.index(docstring.LEAD_IN) < doc.index("Examples") def test_google_flat_content_gets_default_indent() -> None: # malformed section content at the header's own indent: the block # still lands at the default content indent and round-trips - doc = build_docstring(GOOGLE_FLAT_CONTENT_DOC, SOURCE) - lead = next(line for line in doc.splitlines() if line.strip() == LEAD_IN) - assert lead == f" {LEAD_IN}" + doc = docstring.build_docstring(GOOGLE_FLAT_CONTENT_DOC, SOURCE) + lead = next(line for line in doc.splitlines() if line.strip() == docstring.LEAD_IN) + assert lead == f" {docstring.LEAD_IN}" def test_empty_section_body_gets_default_indent() -> None: # an empty Notes section (header, then the next header): our block # still lands inside it at the default content indent - doc = build_docstring("Summary.\n\nNotes:\n\nExamples:\n stuff", SOURCE) + doc = docstring.build_docstring("Summary.\n\nNotes:\n\nExamples:\n stuff", SOURCE) lines = doc.splitlines() notes_idx = lines.index("Notes:") - assert lines[notes_idx + 1] == f" {LEAD_IN}" + assert lines[notes_idx + 1] == f" {docstring.LEAD_IN}" assert doc.index("def clamp(x") < doc.index("Examples:") @@ -170,42 +161,53 @@ def test_empty_section_body_gets_default_indent() -> None: ], ) def test_is_numpy_style(doc: str, expected: bool) -> None: - assert is_numpy_style(doc) is expected + assert docstring.is_numpy_style(doc) is expected def test_extract_no_marker_raises() -> None: with pytest.raises(ValueError, match="no array-vectorize source marker"): - extract_scalar_source("Just a docstring without a marker.") + docstring.extract_scalar_source("Just a docstring without a marker.") def test_extract_empty_block_raises() -> None: with pytest.raises(ValueError, match="no source block"): - extract_scalar_source(f"Summary.\n\n{LEAD_IN}\n\nUnindented tail at column zero.") + docstring.extract_scalar_source( + f"Summary.\n\n{docstring.LEAD_IN}\n\nUnindented tail at column zero." + ) def test_prefixed_summary_accepts_raw_docstrings() -> None: # the fallback wrapper passes func.__doc__ verbatim (leading newline, # function-body indentation) — cleandoc normalizes before prefixing raw = "\n Summary line.\n\n Body.\n " - assert prefixed_summary(raw).splitlines()[0] == f"{DOCSTRING_PREFIX}Summary line." - assert prefixed_summary(None) == DOCSTRING_PREFIX + NO_DOCSTRING_SUMMARY - assert prefixed_summary("") == DOCSTRING_PREFIX + NO_DOCSTRING_SUMMARY + assert ( + docstring.prefixed_summary(raw).splitlines()[0] + == f"{docstring.DOCSTRING_PREFIX}Summary line." + ) + assert ( + docstring.prefixed_summary(None) + == docstring.DOCSTRING_PREFIX + docstring.NO_DOCSTRING_SUMMARY + ) + assert ( + docstring.prefixed_summary("") + == docstring.DOCSTRING_PREFIX + docstring.NO_DOCSTRING_SUMMARY + ) def test_docstring_value_self_check_passes() -> None: - value = _docstring_value(GOOGLE_DOC, SOURCE) - doc = cleandoc(value) - assert doc.splitlines()[0] == f"{DOCSTRING_PREFIX}Clamp ``x``." - assert extract_scalar_source(doc) == SOURCE + value = emitter._docstring_value(GOOGLE_DOC, SOURCE) + doc = inspect.cleandoc(value) + assert doc.splitlines()[0] == f"{docstring.DOCSTRING_PREFIX}Clamp ``x``." + assert docstring.extract_scalar_source(doc) == SOURCE def test_docstring_value_degrades_on_marker_collision() -> None: # the original's docs quote the marker line before our section: the # extraction would grab the quoted block, so the legacy source-only # docstring is emitted instead - quoting_doc = f"Summary.\n\n{LEAD_IN}\n\n quoted block" - value = _docstring_value(quoting_doc, SOURCE) - assert cleandoc(value) == SOURCE + quoting_doc = f"Summary.\n\n{docstring.LEAD_IN}\n\n quoted block" + value = emitter._docstring_value(quoting_doc, SOURCE) + assert inspect.cleandoc(value) == SOURCE def test_docstring_value_degrades_on_tab_indented_source() -> None: @@ -214,6 +216,6 @@ def test_docstring_value_degrades_on_tab_indented_source() -> None: # same pre-existing caveat): degrade to the legacy source-only # docstring, which still embeds the source text verbatim tab_source = "def f(x):\n\treturn x" - value = _docstring_value(None, tab_source) - assert LEAD_IN not in value # legacy format, not the combined one + value = emitter._docstring_value(None, tab_source) + assert docstring.LEAD_IN not in value # legacy format, not the combined one assert "def f(x):\n \treturn x" in value diff --git a/tests/unit/test_extract.py b/tests/unit/test_extract.py index 6809248..cd88e7a 100644 --- a/tests/unit/test_extract.py +++ b/tests/unit/test_extract.py @@ -6,10 +6,10 @@ import numpy as np import pytest -from support import make_module +import support -from array_vectorize.errors import VectorizationError -from array_vectorize.frontend.extract import extract_function +from array_vectorize import errors +from array_vectorize.frontend import extract K = 2.5 ARR = np.asarray([1.0, 2.0]) @@ -22,7 +22,7 @@ def fn_plain(x: float) -> float: def test_plain_function() -> None: - info = extract_function(fn_plain) + info = extract.extract_function(fn_plain) assert info.name == "fn_plain" assert info.params[0].name == "x" assert info.docstring == "Docstring here." @@ -31,7 +31,7 @@ def test_plain_function() -> None: def test_lambda_via_assignment() -> None: - lam = extract_function(lambda x: x + 1.0) # type: ignore[arg-type] + lam = extract.extract_function(lambda x: x + 1.0) # type: ignore[arg-type] assert lam.params[0].name == "x" assert lam.docstring is None @@ -44,7 +44,7 @@ def deco(f): def fn(x: float) -> float: return x - info = extract_function(fn) + info = extract.extract_function(fn) assert info.params[0].name == "x" @@ -55,33 +55,33 @@ def test_unwrap_follows_wrapped() -> None: def wrapper(x: float) -> float: raise AssertionError("never called") - info = extract_function(wrapper) + info = extract.extract_function(wrapper) assert info.name == "fn_plain" def test_resolve_original_follows_marker_chains() -> None: import functools - from array_vectorize import vectorize - from array_vectorize.frontend.extract import resolve_original + import array_vectorize + from array_vectorize.frontend import extract as extract_mod @functools.wraps(fn_plain) def wrapper(x: float) -> float: raise AssertionError("never called") # __wrapped__ chains resolve to the innermost function... - assert resolve_original(wrapper) is fn_plain + assert extract_mod.resolve_original(wrapper) is fn_plain # ...and prior vectorizations chain to the scalar original - vec = vectorize(fn_plain) - assert resolve_original(vec) is fn_plain - assert resolve_original(vectorize(vec)) is fn_plain + vec = array_vectorize.vectorize(fn_plain) + assert extract_mod.resolve_original(vec) is fn_plain + assert extract_mod.resolve_original(array_vectorize.vectorize(vec)) is fn_plain # plain functions pass through unchanged - assert resolve_original(fn_plain) is fn_plain + assert extract_mod.resolve_original(fn_plain) is fn_plain def test_builtin_rejected() -> None: - with pytest.raises(VectorizationError, match="builtin"): - extract_function(math.sqrt) # type: ignore[arg-type] + with pytest.raises(errors.VectorizationError, match="builtin"): + extract.extract_function(math.sqrt) # type: ignore[arg-type] def test_method_rejected() -> None: @@ -89,62 +89,62 @@ class C: def m(self, x: float) -> float: return x - with pytest.raises(VectorizationError, match="plain Python function"): - extract_function(C().m) # type: ignore[arg-type] + with pytest.raises(errors.VectorizationError, match="plain Python function"): + extract.extract_function(C().m) # type: ignore[arg-type] def test_no_source_rejected_with_guidance() -> None: ns: dict[str, object] = {} exec("def f(x): return x", ns) - with pytest.raises(VectorizationError, match="REPL"): - extract_function(ns["f"]) # type: ignore[arg-type] + with pytest.raises(errors.VectorizationError, match="REPL"): + extract.extract_function(ns["f"]) # type: ignore[arg-type] def test_vararg_rejected() -> None: def f(*args: float) -> float: return 0.0 - with pytest.raises(VectorizationError, match=r"\*args"): - extract_function(f) + with pytest.raises(errors.VectorizationError, match=r"\*args"): + extract.extract_function(f) def test_kwarg_rejected() -> None: def f(x: float, **kw: float) -> float: return x - with pytest.raises(VectorizationError, match=r"\*\*kwargs"): - extract_function(f) + with pytest.raises(errors.VectorizationError, match=r"\*\*kwargs"): + extract.extract_function(f) def test_non_literal_default_rejected() -> None: def f(x: float, s: str = "a") -> float: return x - with pytest.raises(VectorizationError, match="literal"): - extract_function(f) + with pytest.raises(errors.VectorizationError, match="literal"): + extract.extract_function(f) def test_none_default_rejected() -> None: def f(x: float, y: float | None = None) -> float: # type: ignore[assignment] return x - with pytest.raises(VectorizationError, match="literal"): - extract_function(f) + with pytest.raises(errors.VectorizationError, match="literal"): + extract.extract_function(f) def test_zero_params_rejected() -> None: def f() -> float: return 1.0 - with pytest.raises(VectorizationError, match="zero-argument"): - extract_function(f) + with pytest.raises(errors.VectorizationError, match="zero-argument"): + extract.extract_function(f) def test_literal_defaults_recorded() -> None: def f(x: float, scale: float = 2.0, n: int = 3, flag: bool = True) -> float: return x * scale - info = extract_function(f) + info = extract.extract_function(f) assert [(p.name, p.default, p.has_default) for p in info.params[1:]] == [ ("scale", 2.0, True), ("n", 3, True), @@ -156,18 +156,20 @@ def test_kwonly_params() -> None: def f(x: float, *, scale: float = 1.0) -> float: return x * scale - info = extract_function(f) + info = extract.extract_function(f) assert info.params[1].kind == "kwonly" assert info.params[1].default == 1.0 def test_from_math_import_resolved_by_identity() -> None: - from math import sqrt + # the from-import is the feature under test: extract must capture the + # locally from-imported math function by identity + from math import sqrt # cleanporter: ignore[CP001] feature under test def f(x: float) -> float: return sqrt(x) # type: ignore[misc] - info = extract_function(f) + info = extract.extract_function(f) assert info.math_funcs == {"sqrt": "sqrt"} @@ -175,7 +177,7 @@ def test_math_module_recorded() -> None: def f(x: float) -> float: return math.exp(x) - info = extract_function(f) + info = extract.extract_function(f) assert info.math_modules == {"math"} @@ -183,7 +185,7 @@ def test_closure_array_captured() -> None: def f(x: float) -> float: return x + ARR[0] * 0.0 + x * 0.0 + x * 1.0 # ARR referenced (subscript rejected later) - info = extract_function(f) + info = extract.extract_function(f) assert info.closure_arrays == {"ARR": ARR} @@ -191,16 +193,16 @@ def test_unresolvable_global_rejected() -> None: def f(x: float) -> float: return x + len(STR_GLOBAL) * 0.0 # type: ignore[misc] - with pytest.raises(VectorizationError, match="unsupported type"): - extract_function(f) + with pytest.raises(errors.VectorizationError, match="unsupported type"): + extract.extract_function(f) def test_unbound_name_rejected() -> None: def f(x: float) -> float: return x + unknown_name # type: ignore[name-defined] # noqa: F821 - with pytest.raises(VectorizationError, match="not resolvable"): - extract_function(f) + with pytest.raises(errors.VectorizationError, match="not resolvable"): + extract.extract_function(f) def test_user_function_captured() -> None: @@ -210,7 +212,7 @@ def helper(y: float) -> float: def f(x: float) -> float: return helper(x) # type: ignore[misc] - info = extract_function(f) + info = extract.extract_function(f) assert info.user_funcs == {"helper": helper} @@ -219,7 +221,7 @@ def f(x: float, x_1: float) -> float: y = x + x_1 return y - info = extract_function(f) + info = extract.extract_function(f) assert {"x", "x_1", "y"} <= info.user_names @@ -229,7 +231,7 @@ def test_np_float64_is_scalar() -> None: def f(x: float) -> float: return x * c - info = extract_function(f) + info = extract.extract_function(f) assert info.closure_scalars == {"c": 1.5} @@ -242,7 +244,7 @@ def test_lambda_probe_fallback_identifies_sibling() -> None: import pathlib import tempfile - from array_vectorize.frontend.lambda_id import _find_target + from array_vectorize.frontend import lambda_id tmp = pathlib.Path(tempfile.mkdtemp()) path = tmp / "probe_siblings.py" @@ -255,12 +257,14 @@ def test_lambda_probe_fallback_identifies_sibling() -> None: for fn, body in ((mod.f1, "lambda x: x + 1"), (mod.f2, "lambda x: x + 2")): tree = ast.parse(inspect.getsource(fn)) - node = _find_target(tree, "", fn, module_source=None) + node = lambda_id._find_target(tree, "", fn, module_source=None) assert ast.unparse(node) == body # a module source that cannot be compiled also falls back to probes tree = ast.parse(inspect.getsource(mod.f2)) - node = _find_target(tree, "", mod.f2, module_source="def broken(:", first_lineno=1) + node = lambda_id._find_target( + tree, "", mod.f2, module_source="def broken(:", first_lineno=1 + ) assert ast.unparse(node) == "lambda x: x + 2" @@ -271,7 +275,7 @@ def test_lambda_probe_fallback_interchangeable() -> None: import pathlib import tempfile - from array_vectorize.frontend.lambda_id import _find_target + from array_vectorize.frontend import lambda_id tmp = pathlib.Path(tempfile.mkdtemp()) path = tmp / "probe_twin.py" @@ -283,7 +287,7 @@ def test_lambda_probe_fallback_interchangeable() -> None: spec.loader.exec_module(mod) tree = ast.parse(inspect.getsource(mod.f2)) - node = _find_target(tree, "", mod.f2, module_source=None) + node = lambda_id._find_target(tree, "", mod.f2, module_source=None) assert ast.unparse(node) == "lambda x: x + 1" @@ -294,7 +298,7 @@ def test_lambda_probe_fallback_rejects_unidentifiable() -> None: import pathlib import tempfile - from array_vectorize.frontend.lambda_id import _find_target + from array_vectorize.frontend import lambda_id tmp = pathlib.Path(tempfile.mkdtemp()) path = tmp / "probe_none.py" @@ -306,31 +310,31 @@ def test_lambda_probe_fallback_rejects_unidentifiable() -> None: spec.loader.exec_module(mod) tree = ast.parse(inspect.getsource(mod.f1)) - with pytest.raises(VectorizationError, match="could not be identified"): - _find_target(tree, "", mod.f3, module_source=None) + with pytest.raises(errors.VectorizationError, match="could not be identified"): + lambda_id._find_target(tree, "", mod.f3, module_source=None) def test_async_function_rejected() -> None: - mod = make_module("async def subject(x):\n return x\n") - with pytest.raises(VectorizationError, match="async"): - extract_function(mod.subject) + mod = support.make_module("async def subject(x):\n return x\n") + with pytest.raises(errors.VectorizationError, match="async"): + extract.extract_function(mod.subject) def test_kwonly_without_default() -> None: - mod = make_module("def subject(x, *, scale):\n return x * scale\n") - info = extract_function(mod.subject) + mod = support.make_module("def subject(x, *, scale):\n return x * scale\n") + info = extract.extract_function(mod.subject) assert info.params[1].kind == "kwonly" assert not info.params[1].has_default def test_expr_lambda_top_level() -> None: - make_module("(lambda x: x + 1)\n") + support.make_module("(lambda x: x + 1)\n") # nothing to grab; just ensure parse path works via direct lambda - info = extract_function(lambda x: x + 1.0) # type: ignore[arg-type] + info = extract.extract_function(lambda x: x + 1.0) # type: ignore[arg-type] assert info.params[0].name == "x" def test_no_target_in_source_rejected() -> None: - mod = make_module("X = 1\n") - with pytest.raises(VectorizationError): - extract_function(mod.X) # type: ignore[arg-type] + mod = support.make_module("X = 1\n") + with pytest.raises(errors.VectorizationError): + extract.extract_function(mod.X) # type: ignore[arg-type] diff --git a/tests/unit/test_fuzz.py b/tests/unit/test_fuzz.py index 869d14a..93da350 100644 --- a/tests/unit/test_fuzz.py +++ b/tests/unit/test_fuzz.py @@ -2,12 +2,12 @@ from __future__ import annotations -from array_vectorize.fuzz import main +from array_vectorize import fuzz def test_fuzzer_smoke() -> None: # deterministic seed; any internal error or value mismatch fails nonzero - assert main(["--seed", "2024", "--cases", "40", "--seconds", "60"]) == 0 + assert fuzz.main(["--seed", "2024", "--cases", "40", "--seconds", "60"]) == 0 def test_fuzzer_detects_mismatch(monkeypatch) -> None: # type: ignore[no-untyped-def] diff --git a/tests/unit/test_ir.py b/tests/unit/test_ir.py index 99cea62..56834b4 100644 --- a/tests/unit/test_ir.py +++ b/tests/unit/test_ir.py @@ -4,120 +4,114 @@ import math -from array_vectorize.ir import ( - Binding, - BinOp, - Call, - Compare, - DType, - Literal, - Logical, - Program, - Ref, - SSAEnv, - UnaryOp, - Where, - is_bool, -) +from array_vectorize import ir class TestNodeBasics: def test_equality_is_structural(self) -> None: - a = BinOp("add", Literal(1, "int"), Ref("x")) - b = BinOp("add", Literal(1, "int"), Ref("x")) - c = BinOp("add", Literal(2, "int"), Ref("x")) + a = ir.BinOp("add", ir.Literal(1, "int"), ir.Ref("x")) + b = ir.BinOp("add", ir.Literal(1, "int"), ir.Ref("x")) + c = ir.BinOp("add", ir.Literal(2, "int"), ir.Ref("x")) assert a == b assert hash(a) == hash(b) assert a != c def test_inf_nan_literals_hashable(self) -> None: - assert hash(Literal(math.inf, "float")) == hash(Literal(math.inf, "float")) - assert hash(Literal(math.nan, "float")) == hash(Literal(math.nan, "float")) + assert hash(ir.Literal(math.inf, "float")) == hash(ir.Literal(math.inf, "float")) + assert hash(ir.Literal(math.nan, "float")) == hash(ir.Literal(math.nan, "float")) def test_nested_nodes_hashable(self) -> None: - p = Program( + p = ir.Program( ("x",), - (Binding("v", Call("sqrt", (Ref("x"),))),), - Where(Compare("lt", Ref("x"), Literal(0, "int")), Literal(0.0, "float"), Ref("v")), + (ir.Binding("v", ir.Call("sqrt", (ir.Ref("x"),))),), + ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int")), + ir.Literal(0.0, "float"), + ir.Ref("v"), + ), ) assert hash(p) == hash( - Program( + ir.Program( ("x",), - (Binding("v", Call("sqrt", (Ref("x"),))),), - Where(Compare("lt", Ref("x"), Literal(0, "int")), Literal(0.0, "float"), Ref("v")), + (ir.Binding("v", ir.Call("sqrt", (ir.Ref("x"),))),), + ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int")), + ir.Literal(0.0, "float"), + ir.Ref("v"), + ), ) ) def test_bool_is_not_int_literal(self) -> None: - assert Literal(True, "bool") != Literal(1, "int") + assert ir.Literal(True, "bool") != ir.Literal(1, "int") class TestIsBool: def test_bool_nodes(self) -> None: - assert is_bool(Compare("lt", Ref("x"), Literal(0, "int"))) - assert is_bool(UnaryOp("not", Ref("x"))) - assert is_bool(Logical("and", (Ref("a"), Ref("b")))) - assert is_bool(Literal(True, "bool")) - assert is_bool(Call("isnan", (Ref("x"),))) - assert is_bool(Call("isinf", (Ref("x"),))) - assert is_bool(Call("isfinite", (Ref("x"),))) + assert ir.is_bool(ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int"))) + assert ir.is_bool(ir.UnaryOp("not", ir.Ref("x"))) + assert ir.is_bool(ir.Logical("and", (ir.Ref("a"), ir.Ref("b")))) + assert ir.is_bool(ir.Literal(True, "bool")) + assert ir.is_bool(ir.Call("isnan", (ir.Ref("x"),))) + assert ir.is_bool(ir.Call("isinf", (ir.Ref("x"),))) + assert ir.is_bool(ir.Call("isfinite", (ir.Ref("x"),))) def test_non_bool_nodes(self) -> None: - assert not is_bool(Ref("x")) - assert not is_bool(Literal(1, "int")) - assert not is_bool(Literal(1.0, "float")) - assert not is_bool(BinOp("add", Ref("x"), Ref("y"))) - assert not is_bool(UnaryOp("neg", Ref("x"))) - assert not is_bool(Call("sqrt", (Ref("x"),))) - assert not is_bool(Where(Ref("c"), Ref("a"), Ref("b"))) - assert not is_bool(DType("int64")) + assert not ir.is_bool(ir.Ref("x")) + assert not ir.is_bool(ir.Literal(1, "int")) + assert not ir.is_bool(ir.Literal(1.0, "float")) + assert not ir.is_bool(ir.BinOp("add", ir.Ref("x"), ir.Ref("y"))) + assert not ir.is_bool(ir.UnaryOp("neg", ir.Ref("x"))) + assert not ir.is_bool(ir.Call("sqrt", (ir.Ref("x"),))) + assert not ir.is_bool(ir.Where(ir.Ref("c"), ir.Ref("a"), ir.Ref("b"))) + assert not ir.is_bool(ir.DType("int64")) class TestSSAEnv: def test_first_binding_keeps_name(self) -> None: - env = SSAEnv(["x"]) + env = ir.SSAEnv(["x"]) assert env.bind("y") == "y" def test_rebind_picks_smallest_free_suffix(self) -> None: - env = SSAEnv(["x"]) + env = ir.SSAEnv(["x"]) env.bind("x") assert env.bind("x") == "x_1" assert env.bind("x") == "x_2" def test_rebind_skips_user_names(self) -> None: # user has variables x and x_1: rebinding x must skip x_1 - env = SSAEnv(["x", "x_1"]) + env = ir.SSAEnv(["x", "x_1"]) env.bind("x") assert env.bind("x") == "x_2" def test_reserved_namespace_name_mangled(self) -> None: # 'xp' is no longer reserved (the namespace var is renamed instead); # names the generated code cannot rename are still mangled - env = SSAEnv(["xp"]) + env = ir.SSAEnv(["xp"]) assert env.bind("xp") == "xp" def test_reserved_hasattr_mangled(self) -> None: - env = SSAEnv(["array_namespace"]) + env = ir.SSAEnv(["array_namespace"]) assert env.bind("array_namespace") == "array_namespace_" - env = SSAEnv(["hasattr"]) + env = ir.SSAEnv(["hasattr"]) assert env.bind("hasattr") == "hasattr_" def test_mangle_skips_taken_mangled_name(self) -> None: # user has both array_namespace and array_namespace_ variables - env = SSAEnv(["array_namespace", "array_namespace_"]) + env = ir.SSAEnv(["array_namespace", "array_namespace_"]) assert env.bind("array_namespace") == "array_namespace_1" def test_fresh_temp_skips_user_names(self) -> None: - env = SSAEnv(["t_1"]) + env = ir.SSAEnv(["t_1"]) assert env.fresh_temp("t") == "t_2" def test_fresh_temp_skips_emitted(self) -> None: - env = SSAEnv([]) + env = ir.SSAEnv([]) env.fresh_temp("t") assert env.fresh_temp("t") == "t_2" def test_reserve_blocks_allocation(self) -> None: - env = SSAEnv(["g"]) + env = ir.SSAEnv(["g"]) env.reserve("g_vec") env.bind("g") # g_vec is reserved and must never be handed out as a binding name diff --git a/tests/unit/test_layering.py b/tests/unit/test_layering.py index d0bfd8c..9473d3d 100644 --- a/tests/unit/test_layering.py +++ b/tests/unit/test_layering.py @@ -9,13 +9,13 @@ from __future__ import annotations import ast +import pathlib import sys -from pathlib import Path import pytest PACKAGE = "array_vectorize" -ROOT = Path(__file__).parents[2] / "src" / PACKAGE +ROOT = pathlib.Path(__file__).parents[2] / "src" / PACKAGE #: the unit of the package root ``__init__.py`` (its stem; a subpackage's #: ``__init__.py`` instead belongs to that subpackage's unit) ROOT_UNIT = "__init__" @@ -26,7 +26,7 @@ DEV_ONLY_MODULES = {"fuzz"} -def _unit(path: Path) -> str: +def _unit(path: pathlib.Path) -> str: """Top-level unit of a source file: a subpackage or a top-level module.""" parts = path.relative_to(ROOT).parts return parts[0] if len(parts) > 1 else path.stem @@ -81,7 +81,7 @@ def _classify(tree: ast.Module, package_dir: list[str]) -> tuple[set[str], set[s return internal, external -def _imports(path: Path) -> tuple[set[str], set[str]]: +def _imports(path: pathlib.Path) -> tuple[set[str], set[str]]: """(package units, external top-level modules) imported by ``path``.""" package_dir = list(path.relative_to(ROOT).parts[:-1]) return _classify(ast.parse(path.read_text()), package_dir) @@ -90,7 +90,7 @@ def _imports(path: Path) -> tuple[set[str], set[str]]: SOURCES = sorted(ROOT.rglob("*.py")) -def _files_of(unit: str) -> list[Path]: +def _files_of(unit: str) -> list[pathlib.Path]: return [path for path in SOURCES if _unit(path) == unit] @@ -144,7 +144,7 @@ def test_known_real_imports() -> None: @pytest.mark.parametrize("path", _files_of("runtime"), ids=lambda p: p.name) -def test_runtime_is_a_stdlib_only_leaf(path: Path) -> None: +def test_runtime_is_a_stdlib_only_leaf(path: pathlib.Path) -> None: # runtime helpers are injected into generated code and run on the # user's backend at call time; they must not drag in the compiler internal, external = _imports(path) @@ -153,13 +153,13 @@ def test_runtime_is_a_stdlib_only_leaf(path: Path) -> None: @pytest.mark.parametrize("path", _files_of("ir"), ids=lambda p: p.name) -def test_ir_depends_on_nothing_in_the_package(path: Path) -> None: +def test_ir_depends_on_nothing_in_the_package(path: pathlib.Path) -> None: internal, _ = _imports(path) assert internal <= {"ir"}, internal @pytest.mark.parametrize("path", SOURCES, ids=lambda p: str(p.relative_to(ROOT))) -def test_only_the_package_root_imports_api(path: Path) -> None: +def test_only_the_package_root_imports_api(path: pathlib.Path) -> None: # the pipeline reaches back into api only through the # helper_vectorizer callback, never by import; the root unit # re-exports api, so a dependency on it reaches api too @@ -175,7 +175,7 @@ def test_only_the_package_root_imports_api(path: Path) -> None: @pytest.mark.parametrize("path", SOURCES, ids=lambda p: str(p.relative_to(ROOT))) -def test_only_runtime_dependencies_are_imported(path: Path) -> None: +def test_only_runtime_dependencies_are_imported(path: pathlib.Path) -> None: if _unit(path) in DEV_ONLY_MODULES: return _, external = _imports(path) diff --git a/tests/unit/test_lower.py b/tests/unit/test_lower.py index ecb46da..ba7bca7 100644 --- a/tests/unit/test_lower.py +++ b/tests/unit/test_lower.py @@ -8,33 +8,19 @@ import numpy as np import pytest -from support import make_module - -from array_vectorize import vectorize -from array_vectorize.errors import VectorizationError -from array_vectorize.frontend.extract import extract_function -from array_vectorize.ir import ( - Binding, - BinOp, - Call, - Compare, - DType, - FuncCall, - Literal, - Logical, - Program, - Ref, - UnaryOp, - Where, -) -from array_vectorize.lower import lower_function +import support + +import array_vectorize +from array_vectorize import errors, ir +from array_vectorize import lower as lower_mod +from array_vectorize.frontend import extract GLOBAL_K = 3 GLOBAL_ARR = np.asarray([10.0, 20.0]) def make_fn(body: str, extra_globals: str = "") -> Callable[..., Any]: - mod = make_module( + mod = support.make_module( "import math\n" "from math import exp\n" "GLOBAL_K = 3\n" @@ -47,43 +33,51 @@ def make_fn(body: str, extra_globals: str = "") -> Callable[..., Any]: return mod.subject -def lower(body: str) -> Program: - return lower_function(extract_function(make_fn(body))).program +def lower(body: str) -> ir.Program: + return lower_mod.lower_function(extract.extract_function(make_fn(body))).program # ----------------------------------------------------------------- expressions def test_const_and_binop() -> None: - assert lower(" return x + 1") == Program( - ("x", "y"), (), BinOp("add", Ref("x"), Literal(1, "int")) + assert lower(" return x + 1") == ir.Program( + ("x", "y"), (), ir.BinOp("add", ir.Ref("x"), ir.Literal(1, "int")) ) def test_arith_ops_map() -> None: p = lower(" return x - y * 2 / 3") - assert p.result == BinOp( - "sub", Ref("x"), BinOp("div", BinOp("mul", Ref("y"), Literal(2, "int")), Literal(3, "int")) + assert p.result == ir.BinOp( + "sub", + ir.Ref("x"), + ir.BinOp("div", ir.BinOp("mul", ir.Ref("y"), ir.Literal(2, "int")), ir.Literal(3, "int")), ) def test_pow_floordiv_mod_bitwise() -> None: - assert lower(" return x ** 2").result == BinOp("pow", Ref("x"), Literal(2, "int")) - assert lower(" return x // 2").result == BinOp("floordiv", Ref("x"), Literal(2, "int")) - assert lower(" return x % 2").result == BinOp("mod", Ref("x"), Literal(2, "int")) - assert lower(" return x & 1").result == BinOp("and", Ref("x"), Literal(1, "int")) - assert lower(" return x | 1").result == BinOp("or", Ref("x"), Literal(1, "int")) - assert lower(" return x ^ 1").result == BinOp("xor", Ref("x"), Literal(1, "int")) - assert lower(" return x << 1").result == BinOp("lshift", Ref("x"), Literal(1, "int")) - assert lower(" return x >> 1").result == BinOp("rshift", Ref("x"), Literal(1, "int")) + assert lower(" return x ** 2").result == ir.BinOp("pow", ir.Ref("x"), ir.Literal(2, "int")) + assert lower(" return x // 2").result == ir.BinOp( + "floordiv", ir.Ref("x"), ir.Literal(2, "int") + ) + assert lower(" return x % 2").result == ir.BinOp("mod", ir.Ref("x"), ir.Literal(2, "int")) + assert lower(" return x & 1").result == ir.BinOp("and", ir.Ref("x"), ir.Literal(1, "int")) + assert lower(" return x | 1").result == ir.BinOp("or", ir.Ref("x"), ir.Literal(1, "int")) + assert lower(" return x ^ 1").result == ir.BinOp("xor", ir.Ref("x"), ir.Literal(1, "int")) + assert lower(" return x << 1").result == ir.BinOp( + "lshift", ir.Ref("x"), ir.Literal(1, "int") + ) + assert lower(" return x >> 1").result == ir.BinOp( + "rshift", ir.Ref("x"), ir.Literal(1, "int") + ) def test_unary_ops() -> None: - assert lower(" return -x").result == UnaryOp("neg", Ref("x")) - assert lower(" return +x").result == UnaryOp("pos", Ref("x")) - assert lower(" return ~x").result == UnaryOp("invert", Ref("x")) - assert lower(" return not (x > 0)").result == UnaryOp( - "not", Compare("gt", Ref("x"), Literal(0, "int")) + assert lower(" return -x").result == ir.UnaryOp("neg", ir.Ref("x")) + assert lower(" return +x").result == ir.UnaryOp("pos", ir.Ref("x")) + assert lower(" return ~x").result == ir.UnaryOp("invert", ir.Ref("x")) + assert lower(" return not (x > 0)").result == ir.UnaryOp( + "not", ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int")) ) @@ -91,100 +85,110 @@ def test_bool_literals_typed() -> None: # bool literal in arithmetic: the whole operation goes through the # runtime _vec_arith helper (bools need a numeric representation; # backends saturate or reject bool arithmetic) - assert lower(" return x + True").result == FuncCall( - "_vec_arith", (Ref("xp"), Literal(1, "int"), Ref("x"), Literal(True, "bool")) + assert lower(" return x + True").result == ir.FuncCall( + "_vec_arith", (ir.Ref("xp"), ir.Literal(1, "int"), ir.Ref("x"), ir.Literal(True, "bool")) ) - assert lower(" return 1.5").result == Literal(1.5, "float") + assert lower(" return 1.5").result == ir.Literal(1.5, "float") def test_math_attribute_call() -> None: - assert lower(" return math.sqrt(x)").result == Call( - "sqrt", (Call("astype", (Call("asarray", (Ref("x"),)), DType("float64"))),) + assert lower(" return math.sqrt(x)").result == ir.Call( + "sqrt", (ir.Call("astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("float64"))),) ) - assert lower(" return math.atan2(x, y)").result == Call( + assert lower(" return math.atan2(x, y)").result == ir.Call( "atan2", ( - Call("astype", (Call("asarray", (Ref("x"),)), DType("float64"))), - Call("astype", (Call("asarray", (Ref("y"),)), DType("float64"))), + ir.Call("astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("float64"))), + ir.Call("astype", (ir.Call("asarray", (ir.Ref("y"),)), ir.DType("float64"))), ), ) - assert lower(" return math.trunc(x)").result == Call( - "astype", (Call("asarray", (Ref("x"),)), DType("int64")) + assert lower(" return math.trunc(x)").result == ir.Call( + "astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("int64")) ) - assert lower(" return math.pow(x, 2)").result == Call( + assert lower(" return math.pow(x, 2)").result == ir.Call( "pow", - (Call("astype", (Call("asarray", (Ref("x"),)), DType("float64"))), Literal(2.0, "float")), + ( + ir.Call("astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("float64"))), + ir.Literal(2.0, "float"), + ), ) def test_from_math_import_call() -> None: - assert lower(" return exp(x)").result == Call( - "exp", (Call("astype", (Call("asarray", (Ref("x"),)), DType("float64"))),) + assert lower(" return exp(x)").result == ir.Call( + "exp", (ir.Call("astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("float64"))),) ) def test_math_constants() -> None: - assert lower(" return math.pi").result == Literal(math.pi, "float") - assert lower(" return math.inf").result == Literal(math.inf, "float") + assert lower(" return math.pi").result == ir.Literal(math.pi, "float") + assert lower(" return math.inf").result == ir.Literal(math.inf, "float") def test_builtins() -> None: - assert lower(" return abs(x)").result == Call("abs", (Call("asarray", (Ref("x"),)),)) - assert lower(" return round(x)").result == Call("round", (Call("asarray", (Ref("x"),)),)) - assert lower(" return min(x, y)").result == FuncCall( + assert lower(" return abs(x)").result == ir.Call( + "abs", (ir.Call("asarray", (ir.Ref("x"),)),) + ) + assert lower(" return round(x)").result == ir.Call( + "round", (ir.Call("asarray", (ir.Ref("x"),)),) + ) + assert lower(" return min(x, y)").result == ir.FuncCall( "_vec_minmax", ( - Ref("xp"), - Literal(True, "bool"), - Call("asarray", (Ref("x"),)), - Call("asarray", (Ref("y"),)), + ir.Ref("xp"), + ir.Literal(True, "bool"), + ir.Call("asarray", (ir.Ref("x"),)), + ir.Call("asarray", (ir.Ref("y"),)), ), ) - assert lower(" return max(x, y, 2.0)").result == FuncCall( + assert lower(" return max(x, y, 2.0)").result == ir.FuncCall( "_vec_minmax", ( - Ref("xp"), - Literal(False, "bool"), - Call("asarray", (Ref("x"),)), - Call("asarray", (Ref("y"),)), - Literal(2.0, "float"), + ir.Ref("xp"), + ir.Literal(False, "bool"), + ir.Call("asarray", (ir.Ref("x"),)), + ir.Call("asarray", (ir.Ref("y"),)), + ir.Literal(2.0, "float"), ), ) - assert lower(" return int(x)").result == Call( - "astype", (Call("asarray", (Ref("x"),)), DType("int64")) + assert lower(" return int(x)").result == ir.Call( + "astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("int64")) ) - assert lower(" return float(x)").result == Call( - "astype", (Call("asarray", (Ref("x"),)), DType("float64")) + assert lower(" return float(x)").result == ir.Call( + "astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("float64")) ) - assert lower(" return bool(x)").result == Call( - "astype", (Call("asarray", (Ref("x"),)), DType("bool")) + assert lower(" return bool(x)").result == ir.Call( + "astype", (ir.Call("asarray", (ir.Ref("x"),)), ir.DType("bool")) ) def test_comparisons_and_chains() -> None: - assert lower(" return x < y").result == Compare("lt", Ref("x"), Ref("y")) - assert lower(" return 0 < x < 1").result == Logical( + assert lower(" return x < y").result == ir.Compare("lt", ir.Ref("x"), ir.Ref("y")) + assert lower(" return 0 < x < 1").result == ir.Logical( "and", ( - Compare("lt", Literal(0, "int"), Ref("x")), - Compare("lt", Ref("x"), Literal(1, "int")), + ir.Compare("lt", ir.Literal(0, "int"), ir.Ref("x")), + ir.Compare("lt", ir.Ref("x"), ir.Literal(1, "int")), ), ) def test_boolop_bool_operands_use_logical() -> None: - assert lower(" return (x < 0) and (y > 0)").result == Logical( + assert lower(" return (x < 0) and (y > 0)").result == ir.Logical( "and", - (Compare("lt", Ref("x"), Literal(0, "int")), Compare("gt", Ref("y"), Literal(0, "int"))), + ( + ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int")), + ir.Compare("gt", ir.Ref("y"), ir.Literal(0, "int")), + ), ) def test_boolop_numeric_value_select() -> None: - assert lower(" return x and y").result == Where( - Compare("ne", Ref("x"), Literal(0, "int")), Ref("y"), Ref("x") + assert lower(" return x and y").result == ir.Where( + ir.Compare("ne", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("y"), ir.Ref("x") ) - assert lower(" return x or y").result == Where( - Compare("ne", Ref("x"), Literal(0, "int")), Ref("x"), Ref("y") + assert lower(" return x or y").result == ir.Where( + ir.Compare("ne", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("x"), ir.Ref("y") ) @@ -193,16 +197,20 @@ def test_boolop_numeric_chain_uses_temp() -> None: assert len(p.bindings) == 1 temp = p.bindings[0] assert temp.name == "v_1" - assert temp.expr == Where(Compare("ne", Ref("x"), Literal(0, "int")), Ref("y"), Ref("x")) - assert p.result == Where(Compare("ne", Ref("v_1"), Literal(0, "int")), Ref("x"), Ref("v_1")) + assert temp.expr == ir.Where( + ir.Compare("ne", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("y"), ir.Ref("x") + ) + assert p.result == ir.Where( + ir.Compare("ne", ir.Ref("v_1"), ir.Literal(0, "int")), ir.Ref("x"), ir.Ref("v_1") + ) def test_ternary_coerces_numeric_condition() -> None: - assert lower(" return x if x else y").result == Where( - Compare("ne", Ref("x"), Literal(0, "int")), Ref("x"), Ref("y") + assert lower(" return x if x else y").result == ir.Where( + ir.Compare("ne", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("x"), ir.Ref("y") ) - assert lower(" return x if x > 0 else y").result == Where( - Compare("gt", Ref("x"), Literal(0, "int")), Ref("x"), Ref("y") + assert lower(" return x if x > 0 else y").result == ir.Where( + ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("x"), ir.Ref("y") ) @@ -212,20 +220,20 @@ def test_ternary_coerces_numeric_condition() -> None: def test_assign_and_rebind() -> None: p = lower(" z = x * 2\n z = z + 1\n return z") assert p.bindings == ( - Binding("z", BinOp("mul", Ref("x"), Literal(2, "int"))), - Binding("z_1", BinOp("add", Ref("z"), Literal(1, "int"))), + ir.Binding("z", ir.BinOp("mul", ir.Ref("x"), ir.Literal(2, "int"))), + ir.Binding("z_1", ir.BinOp("add", ir.Ref("z"), ir.Literal(1, "int"))), ) - assert p.result == Ref("z_1") + assert p.result == ir.Ref("z_1") def test_augassign() -> None: p = lower(" z = x\n z += 1\n z *= 2\n return z") assert p.bindings == ( - Binding("z", Ref("x")), - Binding("z_1", BinOp("add", Ref("z"), Literal(1, "int"))), - Binding("z_2", BinOp("mul", Ref("z_1"), Literal(2, "int"))), + ir.Binding("z", ir.Ref("x")), + ir.Binding("z_1", ir.BinOp("add", ir.Ref("z"), ir.Literal(1, "int"))), + ir.Binding("z_2", ir.BinOp("mul", ir.Ref("z_1"), ir.Literal(2, "int"))), ) - assert p.result == Ref("z_2") + assert p.result == ir.Ref("z_2") def test_rebind_skips_user_names() -> None: @@ -235,13 +243,15 @@ def test_rebind_skips_user_names() -> None: def test_closure_scalar_frozen() -> None: - assert lower(" return x * GLOBAL_K").result == BinOp("mul", Ref("x"), Literal(3, "int")) + assert lower(" return x * GLOBAL_K").result == ir.BinOp( + "mul", ir.Ref("x"), ir.Literal(3, "int") + ) def test_closure_array_hidden_param() -> None: - lf = lower_function(extract_function(make_fn(" return x + GLOBAL_ARR"))) + lf = lower_mod.lower_function(extract.extract_function(make_fn(" return x + GLOBAL_ARR"))) assert lf.program.params == ("x", "y", "GLOBAL_ARR") - assert lf.program.result == BinOp("add", Ref("x"), Ref("GLOBAL_ARR")) + assert lf.program.result == ir.BinOp("add", ir.Ref("x"), ir.Ref("GLOBAL_ARR")) assert len(lf.hidden_params) == 1 name, default = lf.hidden_params[0] assert name == "GLOBAL_ARR" @@ -251,11 +261,11 @@ def test_closure_array_hidden_param() -> None: def test_reserved_param_name_kept_namespace_renamed() -> None: # a parameter named 'xp' keeps its name (keyword calls work); the # generated namespace variable renames itself instead - mod = make_module("def subject(xp):\n return xp + 1\n") - lf = lower_function(extract_function(mod.subject)) + mod = support.make_module("def subject(xp):\n return xp + 1\n") + lf = lower_mod.lower_function(extract.extract_function(mod.subject)) assert lf.param_names == ["xp"] assert lf.namespace_var == "xp_1" - assert lf.program.result == BinOp("add", Ref("xp"), Literal(1, "int")) + assert lf.program.result == ir.BinOp("add", ir.Ref("xp"), ir.Literal(1, "int")) # ----------------------------------------------------------------- rejections @@ -272,23 +282,23 @@ def test_reserved_param_name_kept_namespace_renamed() -> None: ], ) def test_expression_rejections(body: str) -> None: - with pytest.raises(VectorizationError): + with pytest.raises(errors.VectorizationError): lower(body) def test_reference_before_assignment() -> None: - with pytest.raises(VectorizationError, match="before assignment"): + with pytest.raises(errors.VectorizationError, match="before assignment"): lower(" z = z + 1\n return z") def test_closure_scalar_read_after_local_assign_rejected() -> None: # GLOBAL_K is assigned locally -> reads before the assignment are unbound - with pytest.raises(VectorizationError, match="before assignment"): + with pytest.raises(errors.VectorizationError, match="before assignment"): lower(" z = GLOBAL_K + 1\n GLOBAL_K = 2\n return z") def test_calling_variable_rejected() -> None: - with pytest.raises(VectorizationError, match="calling variable"): + with pytest.raises(errors.VectorizationError, match="calling variable"): lower(" f = x\n return f(x)") @@ -298,28 +308,29 @@ def test_calling_variable_rejected() -> None: def test_if_else_merge() -> None: p = lower(" if x > 0:\n r = x\n else:\n r = 0.0\n return r") assert [b.name for b in p.bindings] == ["r_1", "r_2", "r"] - assert p.bindings[0] == Binding("r_1", Ref("x")) - assert p.bindings[1] == Binding("r_2", Literal(0.0, "float")) - assert p.bindings[2] == Binding( - "r", Where(Compare("gt", Ref("x"), Literal(0, "int")), Ref("r_1"), Ref("r_2")) + assert p.bindings[0] == ir.Binding("r_1", ir.Ref("x")) + assert p.bindings[1] == ir.Binding("r_2", ir.Literal(0.0, "float")) + assert p.bindings[2] == ir.Binding( + "r", + ir.Where(ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("r_1"), ir.Ref("r_2")), ) - assert p.result == Ref("r") + assert p.result == ir.Ref("r") def test_if_else_merge_with_prior_binding() -> None: p = lower(" r = 0.0\n if x > 0:\n r = x\n else:\n r = 1.0\n return r") assert [b.name for b in p.bindings] == ["r", "r_1", "r_2", "r_3"] - assert p.result == Ref("r_3") + assert p.result == ir.Ref("r_3") def test_if_without_else_and_prior_binding() -> None: p = lower(" r = 0.0\n if x > 0:\n r = x\n return r") # then binds r_1; else keeps r; merge where(cond, r_1, r) assert [b.name for b in p.bindings] == ["r", "r_1", "r_2"] - assert p.bindings[2].expr == Where( - Compare("gt", Ref("x"), Literal(0, "int")), Ref("r_1"), Ref("r") + assert p.bindings[2].expr == ir.Where( + ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int")), ir.Ref("r_1"), ir.Ref("r") ) - assert p.result == Ref("r_2") + assert p.result == ir.Ref("r_2") def test_elif_chain() -> None: @@ -334,21 +345,29 @@ def test_elif_chain() -> None: ) names = [b.name for b in p.bindings] assert names == ["r_1", "r_2", "r_3", "r", "r_4"] - assert p.result == Ref("r_4") + assert p.result == ir.Ref("r_4") def test_early_return_top_level() -> None: p = lower(" if x < 0:\n return 0.0\n return x * math.exp(-x)") assert p.bindings == () - assert p.result == Where( - Compare("lt", Ref("x"), Literal(0, "int")), - Literal(0.0, "float"), - BinOp( + assert p.result == ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int")), + ir.Literal(0.0, "float"), + ir.BinOp( "mul", - Ref("x"), - Call( + ir.Ref("x"), + ir.Call( "exp", - (Call("astype", (Call("asarray", (UnaryOp("neg", Ref("x")),)), DType("float64"))),), + ( + ir.Call( + "astype", + ( + ir.Call("asarray", (ir.UnaryOp("neg", ir.Ref("x")),)), + ir.DType("float64"), + ), + ), + ), ), ), ) @@ -358,11 +377,13 @@ def test_sequential_same_condition_first_return_wins() -> None: p = lower( " if x < 0:\n return 1.0\n if x < 0:\n return 2.0\n return 3.0" ) - inner = Where( - Compare("lt", Ref("x"), Literal(0, "int")), Literal(2.0, "float"), Literal(3.0, "float") + inner = ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int")), + ir.Literal(2.0, "float"), + ir.Literal(3.0, "float"), ) - assert p.result == Where( - Compare("lt", Ref("x"), Literal(0, "int")), Literal(1.0, "float"), inner + assert p.result == ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int")), ir.Literal(1.0, "float"), inner ) @@ -377,12 +398,18 @@ def test_nested_both_return_with_outer_pending() -> None: " return 2.0\n" " return 3.0" ) - inner = Where( - Compare("gt", Ref("x"), Literal(2, "int")), Literal(1.0, "float"), Literal(2.0, "float") + inner = ir.Where( + ir.Compare("gt", ir.Ref("x"), ir.Literal(2, "int")), + ir.Literal(1.0, "float"), + ir.Literal(2.0, "float"), + ) + middle = ir.Where( + ir.Compare("gt", ir.Ref("x"), ir.Literal(1, "int")), inner, ir.Literal(3.0, "float") ) - middle = Where(Compare("gt", Ref("x"), Literal(1, "int")), inner, Literal(3.0, "float")) - assert p.result == Where( - Compare("lt", Ref("x"), UnaryOp("neg", Literal(5, "int"))), Literal(9.0, "float"), middle + assert p.result == ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.UnaryOp("neg", ir.Literal(5, "int"))), + ir.Literal(9.0, "float"), + middle, ) @@ -395,29 +422,29 @@ def test_then_return_with_inner_else_pending() -> None: " return 2.0\n" " return 3.0" ) - inner = Where( - Compare("lt", Ref("x"), UnaryOp("neg", Literal(1, "int"))), - Literal(2.0, "float"), - Literal(3.0, "float"), + inner = ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.UnaryOp("neg", ir.Literal(1, "int"))), + ir.Literal(2.0, "float"), + ir.Literal(3.0, "float"), ) - assert p.result == Where( - Compare("gt", Ref("x"), Literal(0, "int")), Literal(1.0, "float"), inner + assert p.result == ir.Where( + ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int")), ir.Literal(1.0, "float"), inner ) def test_early_return_in_else_branch() -> None: p = lower(" if x < 0:\n y = -x\n else:\n return 0.0\n return y") - assert p.result == Where( - UnaryOp("not", Compare("lt", Ref("x"), Literal(0, "int"))), - Literal(0.0, "float"), - Ref("y_1"), + assert p.result == ir.Where( + ir.UnaryOp("not", ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int"))), + ir.Literal(0.0, "float"), + ir.Ref("y_1"), ) def test_both_branches_return() -> None: p = lower(" if x < 0:\n return 0.0\n else:\n return x\n") - assert p.result == Where( - Compare("lt", Ref("x"), Literal(0, "int")), Literal(0.0, "float"), Ref("x") + assert p.result == ir.Where( + ir.Compare("lt", ir.Ref("x"), ir.Literal(0, "int")), ir.Literal(0.0, "float"), ir.Ref("x") ) @@ -429,14 +456,14 @@ def test_nested_if_with_returns() -> None: " return 2.0\n" " return 3.0" ) - assert p.result == Where( - Compare("gt", Ref("x"), Literal(0, "int")), - Where( - Compare("gt", Ref("x"), Literal(10, "int")), - Literal(1.0, "float"), - Literal(2.0, "float"), + assert p.result == ir.Where( + ir.Compare("gt", ir.Ref("x"), ir.Literal(0, "int")), + ir.Where( + ir.Compare("gt", ir.Ref("x"), ir.Literal(10, "int")), + ir.Literal(1.0, "float"), + ir.Literal(2.0, "float"), ), - Literal(3.0, "float"), + ir.Literal(3.0, "float"), ) @@ -452,38 +479,38 @@ def test_pending_return_qualified_by_branch() -> None: " y = 3.0\n" " return y" ) - assert isinstance(p.result, Where) + assert isinstance(p.result, ir.Where) outer_cond = p.result.cond - assert isinstance(outer_cond, Logical) + assert isinstance(outer_cond, ir.Logical) assert outer_cond.op == "and" def test_bare_truthiness_if_rejected() -> None: - with pytest.raises(VectorizationError, match="truthiness"): + with pytest.raises(errors.VectorizationError, match="truthiness"): lower(" if x:\n return x\n return 0.0") def test_maybe_unbound_read_rejected() -> None: - with pytest.raises(VectorizationError, match="may be unbound"): + with pytest.raises(errors.VectorizationError, match="may be unbound"): lower(" if x > 0:\n z = 1.0\n return z") def test_maybe_unbound_not_read_is_fine() -> None: p = lower(" if x > 0:\n z = 1.0\n return x") - assert p.result == Ref("x") + assert p.result == ir.Ref("x") # z_1 is dead; DCE (in _optimize) drops it later assert [b.name for b in p.bindings] == ["z_1"] def test_not_all_paths_return_rejected() -> None: - with pytest.raises(VectorizationError, match="paths return"): + with pytest.raises(errors.VectorizationError, match="paths return"): lower(" if x > 0:\n return x") def test_lambda_lowering() -> None: - mod = make_module("subject = lambda x: x + 1.0\n") - lf = lower_function(extract_function(mod.subject)) - assert lf.program.result == BinOp("add", Ref("x"), Literal(1.0, "float")) + mod = support.make_module("subject = lambda x: x + 1.0\n") + lf = lower_mod.lower_function(extract.extract_function(mod.subject)) + assert lf.program.result == ir.BinOp("add", ir.Ref("x"), ir.Literal(1.0, "float")) def test_kinds_restore_revokes_post_snapshot_facts() -> None: @@ -491,9 +518,9 @@ def test_kinds_restore_revokes_post_snapshot_facts() -> None: # pass: restore REPLACES the kind facts (old dict-reassignment # semantics), it must not leave a present-None entry that blocks # setdefault - from array_vectorize.lower.kinds import Kinds + from array_vectorize.lower import kinds - k = Kinds() + k = kinds.Kinds() snap = k.snapshot_facts() # loop name not yet established k.setdefault_kind("i", "int") # post-snapshot insert (loop-var default) k.restore_facts(snap) # widening re-lower: state rolls back @@ -506,9 +533,9 @@ def test_kinds_restore_preserves_present_none_kind() -> None: # a pre-bound loop var with unknown kind establishes a present-None # fact; the restore must keep both presence and value so the later # "int" default stays blocked (old dict setdefault semantics) - from array_vectorize.lower.kinds import Kinds + from array_vectorize.lower import kinds - k = Kinds() + k = kinds.Kinds() k.set_kind("i", None) # phi with unknown kind snap = k.snapshot_facts() k.set_kind("i", "int") # discarded pass wrote a kind @@ -522,30 +549,30 @@ def test_kinds_restore_preserves_present_none_kind() -> None: def test_calling_math_module_rejected() -> None: - with pytest.raises(VectorizationError, match="math module"): - vectorize(make_fn(" return math(x)")) + with pytest.raises(errors.VectorizationError, match="math module"): + array_vectorize.vectorize(make_fn(" return math(x)")) def test_math_const_called_rejected() -> None: - with pytest.raises(VectorizationError, match="cannot be called"): - vectorize(make_fn(" return math.pi(x)")) + with pytest.raises(errors.VectorizationError, match="cannot be called"): + array_vectorize.vectorize(make_fn(" return math.pi(x)")) def test_calling_shadowed_builtin_rejected() -> None: - with pytest.raises(VectorizationError, match="calling variable"): - vectorize(make_fn(" abs = x\n return abs(x)")) + with pytest.raises(errors.VectorizationError, match="calling variable"): + array_vectorize.vectorize(make_fn(" abs = x\n return abs(x)")) def test_not_all_paths_return_via_branch_fallthrough() -> None: - with pytest.raises(VectorizationError, match="paths return"): - vectorize( + with pytest.raises(errors.VectorizationError, match="paths return"): + array_vectorize.vectorize( make_fn(" if x > 0:\n return 1.0\n if x < -1:\n return 2.0\n") ) def test_maybe_unbound_after_nested_branch() -> None: - with pytest.raises(VectorizationError, match="may be unbound"): - vectorize( + with pytest.raises(errors.VectorizationError, match="may be unbound"): + array_vectorize.vectorize( make_fn( " if x > 0:\n" " if x > 1:\n" diff --git a/tests/unit/test_optimize.py b/tests/unit/test_optimize.py index e204574..4839b07 100644 --- a/tests/unit/test_optimize.py +++ b/tests/unit/test_optimize.py @@ -6,136 +6,135 @@ import pytest -from array_vectorize.ir import ( - Binding, - BinOp, - Call, - Compare, - Literal, - Program, - Ref, - SSAEnv, - UnaryOp, - Where, -) -from array_vectorize.optimize import const_fold, cse, dce, optimize -from array_vectorize.optimize import dce as dce_pass +from array_vectorize import ir +from array_vectorize import optimize as optimize_mod + +# optimize's documented submodule/function name collision (the NOTE in +# optimize/__init__.py): the from-import is the supported spelling +from array_vectorize.optimize import cse, dce # cleanporter: ignore[CP002] name collision +from array_vectorize.optimize import dce as dce_pass # cleanporter: ignore[CP002] name collision -def lit(v: float | int, kind: str = "auto") -> Literal: +def lit(v: float | int, kind: str = "auto") -> ir.Literal: if kind == "auto": kind = "bool" if isinstance(v, bool) else "int" if isinstance(v, int) else "float" - return Literal(v, kind) + return ir.Literal(v, kind) # ---------------------------------------------------------------- const-fold def test_fold_arithmetic() -> None: - folded = const_fold(Program((), (), BinOp("add", lit(2), lit(3)))) + folded = optimize_mod.const_fold(ir.Program((), (), ir.BinOp("add", lit(2), lit(3)))) assert folded.result == lit(5) def test_fold_nested() -> None: - expr = BinOp("mul", BinOp("add", lit(2), lit(3)), lit(4)) - assert const_fold(Program((), (), expr)).result == lit(20) + expr = ir.BinOp("mul", ir.BinOp("add", lit(2), lit(3)), lit(4)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == lit(20) def test_fold_div_makes_float() -> None: - assert const_fold(Program((), (), BinOp("div", lit(6), lit(3)))).result == lit(2.0, "float") + assert optimize_mod.const_fold( + ir.Program((), (), ir.BinOp("div", lit(6), lit(3))) + ).result == lit(2.0, "float") def test_fold_div_by_zero_skipped() -> None: - expr = BinOp("div", lit(1.0), lit(0.0)) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.BinOp("div", lit(1.0), lit(0.0)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr def test_fold_mod_by_zero_skipped() -> None: - expr = BinOp("mod", lit(1.0), lit(0.0)) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.BinOp("mod", lit(1.0), lit(0.0)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr def test_fold_pow_skipped() -> None: - expr = BinOp("pow", lit(2), lit(3)) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.BinOp("pow", lit(2), lit(3)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr def test_fold_int64_overflow_skipped() -> None: - expr = BinOp("mul", lit(2**62), lit(4)) - assert const_fold(Program((), (), expr)).result == expr - expr = UnaryOp("neg", lit(-(2**63))) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.BinOp("mul", lit(2**62), lit(4)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr + expr = ir.UnaryOp("neg", lit(-(2**63))) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr def test_fold_bitwise_on_bools_keeps_bool() -> None: - assert const_fold(Program((), (), BinOp("and", lit(True), lit(False)))).result == lit( - False, "bool" - ) + assert optimize_mod.const_fold( + ir.Program((), (), ir.BinOp("and", lit(True), lit(False))) + ).result == lit(False, "bool") def test_fold_unary() -> None: - assert const_fold(Program((), (), UnaryOp("neg", lit(5)))).result == lit(-5) - assert const_fold(Program((), (), UnaryOp("invert", lit(5)))).result == lit(-6) - assert const_fold(Program((), (), UnaryOp("pos", lit(5)))).result == lit(5) + assert optimize_mod.const_fold(ir.Program((), (), ir.UnaryOp("neg", lit(5)))).result == lit(-5) + assert optimize_mod.const_fold(ir.Program((), (), ir.UnaryOp("invert", lit(5)))).result == lit( + -6 + ) + assert optimize_mod.const_fold(ir.Program((), (), ir.UnaryOp("pos", lit(5)))).result == lit(5) def test_fold_not_is_exact_for_nan() -> None: - assert const_fold(Program((), (), UnaryOp("not", lit(math.nan, "float")))).result == lit( - False, "bool" - ) - assert const_fold(Program((), (), UnaryOp("not", lit(0.0, "float")))).result == lit( - True, "bool" - ) + assert optimize_mod.const_fold( + ir.Program((), (), ir.UnaryOp("not", lit(math.nan, "float"))) + ).result == lit(False, "bool") + assert optimize_mod.const_fold( + ir.Program((), (), ir.UnaryOp("not", lit(0.0, "float"))) + ).result == lit(True, "bool") def test_fold_pure_calls_over_literals() -> None: # pure xp functions with literal args fold (keeps plain scalars out of # generated code: strict backends require arrays) - assert const_fold(Program((), (), Call("sqrt", (lit(4),)))).result == lit(2.0) - expr = Call("sqrt", (lit(-1),)) # math.sqrt(-1) raises: no fold - assert const_fold(Program((), (), expr)).result == expr - expr = Call("where", (lit(1), lit(2), lit(3))) # not in the fold table - assert const_fold(Program((), (), expr)).result == expr - expr = Call("sqrt", (Ref("x"),)) # non-literal arg: no fold - assert const_fold(Program((), (), expr)).result == expr + assert optimize_mod.const_fold(ir.Program((), (), ir.Call("sqrt", (lit(4),)))).result == lit( + 2.0 + ) + expr = ir.Call("sqrt", (lit(-1),)) # math.sqrt(-1) raises: no fold + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr + expr = ir.Call("where", (lit(1), lit(2), lit(3))) # not in the fold table + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr + expr = ir.Call("sqrt", (ir.Ref("x"),)) # non-literal arg: no fold + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr # ----------------------------------------------------------------------- DCE def test_dce_drops_unreachable_bindings() -> None: - p = Program( + p = ir.Program( ("x",), ( - Binding("a", Call("sqrt", (Ref("x"),))), - Binding("dead", Call("exp", (Ref("x"),))), + ir.Binding("a", ir.Call("sqrt", (ir.Ref("x"),))), + ir.Binding("dead", ir.Call("exp", (ir.Ref("x"),))), ), - Ref("a"), + ir.Ref("a"), ) - assert dce(p).bindings == (Binding("a", Call("sqrt", (Ref("x"),))),) + assert dce(p).bindings == (ir.Binding("a", ir.Call("sqrt", (ir.Ref("x"),))),) def test_dce_keeps_chains() -> None: - p = Program( + p = ir.Program( ("x",), ( - Binding("a", Call("sqrt", (Ref("x"),))), - Binding("b", Call("exp", (Ref("a"),))), - Binding("dead", Ref("x")), + ir.Binding("a", ir.Call("sqrt", (ir.Ref("x"),))), + ir.Binding("b", ir.Call("exp", (ir.Ref("a"),))), + ir.Binding("dead", ir.Ref("x")), ), - Ref("b"), + ir.Ref("b"), ) assert [b.name for b in dce(p).bindings] == ["a", "b"] def test_dce_single_reverse_pass_handles_late_uses() -> None: - p = Program( + p = ir.Program( ("x",), ( - Binding("a", Ref("x")), - Binding("b", Call("sqrt", (Ref("a"),))), + ir.Binding("a", ir.Ref("x")), + ir.Binding("b", ir.Call("sqrt", (ir.Ref("a"),))), ), - Ref("b"), + ir.Ref("b"), ) assert [b.name for b in dce(p).bindings] == ["a", "b"] @@ -144,58 +143,58 @@ def test_dce_single_reverse_pass_handles_late_uses() -> None: def test_cse_hoists_duplicated_call() -> None: - inner = Call("sqrt", (Ref("x"),)) - p = Program( + inner = ir.Call("sqrt", (ir.Ref("x"),)) + p = ir.Program( ("x",), - (Binding("a", BinOp("add", inner, inner)),), - BinOp("mul", inner, Ref("a")), + (ir.Binding("a", ir.BinOp("add", inner, inner)),), + ir.BinOp("mul", inner, ir.Ref("a")), ) - out = cse(p, SSAEnv(set())) + out = cse(p, ir.SSAEnv(set())) assert out.bindings[0].name == "t_1" assert out.bindings[0].expr == inner - assert out.bindings[1] == Binding("a", BinOp("add", Ref("t_1"), Ref("t_1"))) - assert out.result == BinOp("mul", Ref("t_1"), Ref("a")) + assert out.bindings[1] == ir.Binding("a", ir.BinOp("add", ir.Ref("t_1"), ir.Ref("t_1"))) + assert out.result == ir.BinOp("mul", ir.Ref("t_1"), ir.Ref("a")) def test_cse_no_duplicates_unchanged() -> None: - p = Program( + p = ir.Program( ("x",), - (Binding("a", Call("sqrt", (Ref("x"),))),), - Call("sqrt", (Ref("a"),)), + (ir.Binding("a", ir.Call("sqrt", (ir.Ref("x"),))),), + ir.Call("sqrt", (ir.Ref("a"),)), ) - assert cse(p, SSAEnv(set())) == p + assert cse(p, ir.SSAEnv(set())) == p def test_cse_where_eligible_literals_not() -> None: - where = Where(Compare("gt", Ref("x"), lit(0)), lit(1.0, "float"), lit(2.0, "float")) - p = Program(("x",), (), BinOp("add", where, where)) - out = cse(p, SSAEnv(set())) - assert out.bindings == (Binding("t_1", where),) - assert out.result == BinOp("add", Ref("t_1"), Ref("t_1")) + where = ir.Where(ir.Compare("gt", ir.Ref("x"), lit(0)), lit(1.0, "float"), lit(2.0, "float")) + p = ir.Program(("x",), (), ir.BinOp("add", where, where)) + out = cse(p, ir.SSAEnv(set())) + assert out.bindings == (ir.Binding("t_1", where),) + assert out.result == ir.BinOp("add", ir.Ref("t_1"), ir.Ref("t_1")) def test_cse_nested_fixpoint() -> None: - inner = Call("sqrt", (Ref("x"),)) - outer1 = Where(Compare("gt", inner, lit(0)), inner, lit(0.0, "float")) - p = Program(("x",), (), BinOp("add", outer1, outer1)) - out = cse(p, SSAEnv(set())) + inner = ir.Call("sqrt", (ir.Ref("x"),)) + outer1 = ir.Where(ir.Compare("gt", inner, lit(0)), inner, lit(0.0, "float")) + p = ir.Program(("x",), (), ir.BinOp("add", outer1, outer1)) + out = cse(p, ir.SSAEnv(set())) names = [b.name for b in out.bindings] # inner temp created first, then (next pass) the rebuilt outer assert names == ["t_1", "t_2"] - assert out.result == BinOp("add", Ref("t_2"), Ref("t_2")) + assert out.result == ir.BinOp("add", ir.Ref("t_2"), ir.Ref("t_2")) def test_cse_temp_inserted_before_first_user() -> None: - inner = Call("sqrt", (Ref("x"),)) - p = Program( + inner = ir.Call("sqrt", (ir.Ref("x"),)) + p = ir.Program( ("x",), ( - Binding("a", Ref("x")), - Binding("b", BinOp("add", inner, inner)), + ir.Binding("a", ir.Ref("x")), + ir.Binding("b", ir.BinOp("add", inner, inner)), ), - Ref("b"), + ir.Ref("b"), ) - out = cse(p, SSAEnv(set())) + out = cse(p, ir.SSAEnv(set())) assert [b.name for b in out.bindings] == ["a", "t_1", "b"] @@ -204,21 +203,21 @@ def test_cse_temp_inserted_before_first_user() -> None: def test_optimize_pipeline() -> None: # (x + (2*3)) duplicated, dead binding present - sub = BinOp("add", Ref("x"), BinOp("mul", lit(2), lit(3))) - p = Program( + sub = ir.BinOp("add", ir.Ref("x"), ir.BinOp("mul", lit(2), lit(3))) + p = ir.Program( ("x",), ( - Binding("dead", Call("exp", (Ref("x"),))), - Binding("a", Call("abs", (sub,))), + ir.Binding("dead", ir.Call("exp", (ir.Ref("x"),))), + ir.Binding("a", ir.Call("abs", (sub,))), ), - BinOp("add", Call("abs", (sub,)), Ref("a")), + ir.BinOp("add", ir.Call("abs", (sub,)), ir.Ref("a")), ) - out = optimize(p, {"x", "dead", "a"}) - folded_sub = BinOp("add", Ref("x"), lit(6)) + out = optimize_mod.optimize(p, {"x", "dead", "a"}) + folded_sub = ir.BinOp("add", ir.Ref("x"), lit(6)) assert out.bindings[0].name == "t_1" - assert out.bindings[0].expr == Call("abs", (folded_sub,)) - assert out.bindings[1] == Binding("a", Ref("t_1")) - assert out.result == BinOp("add", Ref("t_1"), Ref("a")) + assert out.bindings[0].expr == ir.Call("abs", (folded_sub,)) + assert out.bindings[1] == ir.Binding("a", ir.Ref("t_1")) + assert out.result == ir.BinOp("add", ir.Ref("t_1"), ir.Ref("a")) # ------------------------------------------------------- const-fold branches @@ -239,51 +238,57 @@ def test_optimize_pipeline() -> None: ], ) def test_fold_binop_ops(op: str, a: float, b: float, expected: float) -> None: - out = const_fold(Program((), (), BinOp(op, lit(a), lit(b)))) + out = optimize_mod.const_fold(ir.Program((), (), ir.BinOp(op, lit(a), lit(b)))) assert out.result == lit(expected) def test_fold_unary_ops_all() -> None: - assert const_fold(Program((), (), UnaryOp("invert", lit(0)))).result == lit(-1) - assert const_fold(Program((), (), UnaryOp("pos", lit(-2)))).result == lit(-2) - assert const_fold(Program((), (), UnaryOp("neg", lit(2.5)))).result == lit(-2.5) - assert const_fold(Program((), (), UnaryOp("not", lit(True)))).result == lit(False) + assert optimize_mod.const_fold(ir.Program((), (), ir.UnaryOp("invert", lit(0)))).result == lit( + -1 + ) + assert optimize_mod.const_fold(ir.Program((), (), ir.UnaryOp("pos", lit(-2)))).result == lit(-2) + assert optimize_mod.const_fold(ir.Program((), (), ir.UnaryOp("neg", lit(2.5)))).result == lit( + -2.5 + ) + assert optimize_mod.const_fold(ir.Program((), (), ir.UnaryOp("not", lit(True)))).result == lit( + False + ) def test_fold_skips_bad_ops() -> None: # float bitwise -> TypeError at runtime -> no fold - expr = BinOp("and", lit(1.5), lit(2.5)) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.BinOp("and", lit(1.5), lit(2.5)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr # unary invert on float -> no fold - expr = UnaryOp("invert", lit(1.5)) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.UnaryOp("invert", lit(1.5)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr # unknown unary op passes through - expr = UnaryOp("weird", lit(1)) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.UnaryOp("weird", lit(1)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr # floordiv by zero - expr = BinOp("floordiv", lit(1), lit(0)) - assert const_fold(Program((), (), expr)).result == expr + expr = ir.BinOp("floordiv", lit(1), lit(0)) + assert optimize_mod.const_fold(ir.Program((), (), expr)).result == expr def test_fold_zero_times_infinity() -> None: # inf * 0 -> nan: exact IEEE fold - out = const_fold(Program((), (), BinOp("mul", lit(math.inf), lit(0.0)))) + out = optimize_mod.const_fold(ir.Program((), (), ir.BinOp("mul", lit(math.inf), lit(0.0)))) assert math.isnan(out.result.value) # type: ignore[attr-defined] def test_dce_keeps_result_only_binding() -> None: - p = Program(("x",), (Binding("a", Call("abs", (Ref("x"),))),), Ref("a")) + p = ir.Program(("x",), (ir.Binding("a", ir.Call("abs", (ir.Ref("x"),))),), ir.Ref("a")) assert dce_pass(p).bindings == p.bindings def test_cse_stops_at_fixpoint() -> None: - inner = Call("sqrt", (Ref("x"),)) - p = Program(("x",), (), BinOp("add", inner, inner)) - out = cse(p, SSAEnv(set())) - assert out.bindings == (Binding("t_1", inner),) - assert out.result == BinOp("add", Ref("t_1"), Ref("t_1")) + inner = ir.Call("sqrt", (ir.Ref("x"),)) + p = ir.Program(("x",), (), ir.BinOp("add", inner, inner)) + out = cse(p, ir.SSAEnv(set())) + assert out.bindings == (ir.Binding("t_1", inner),) + assert out.result == ir.BinOp("add", ir.Ref("t_1"), ir.Ref("t_1")) def test_optimize_no_user_names() -> None: - p = Program(("x",), (Binding("a", Ref("x")),), Ref("a")) - assert optimize(p).result == Ref("a") + p = ir.Program(("x",), (ir.Binding("a", ir.Ref("x")),), ir.Ref("a")) + assert optimize_mod.optimize(p).result == ir.Ref("a") diff --git a/tests/unit/test_runtime.py b/tests/unit/test_runtime.py index cd86e45..c90c79a 100644 --- a/tests/unit/test_runtime.py +++ b/tests/unit/test_runtime.py @@ -10,9 +10,7 @@ import numpy as np -from array_vectorize.runtime.arith import ArithOp, _vec_arith -from array_vectorize.runtime.dtype import _describe_dtype, _fits_dtype -from array_vectorize.runtime.minmax import _vec_minmax +from array_vectorize.runtime import arith, dtype, minmax # the array namespace passed to the runtime helpers, as generated code # binds it @@ -29,7 +27,7 @@ def test_vec_arith_bool_add_is_exact_int64() -> None: # Python: True + True == 2, True + False == 1, False + True == 1 a = np.asarray([True, True, False]) b = np.asarray([True, False, True]) - out = _vec_arith(xp, int(ArithOp.ADD), a, b) + out = arith._vec_arith(xp, int(arith.ArithOp.ADD), a, b) assert out.dtype == np.dtype("int64") assert np.array_equal(out, [2, 1, 1]) @@ -39,7 +37,7 @@ def test_vec_arith_uint64_add_bool_is_modular_uint64() -> None: # modular past 2**64 (documented backend behavior): (2**64-1) + 1 == 0 a = np.asarray([3, U64_MAX], dtype=np.uint64) b = np.asarray([True, True]) - out = _vec_arith(xp, int(ArithOp.ADD), a, b) + out = arith._vec_arith(xp, int(arith.ArithOp.ADD), a, b) assert out.dtype == np.dtype("uint64") assert np.array_equal(out, [4, 0]) @@ -48,7 +46,7 @@ def test_vec_arith_uint64_sub_crossing_zero_is_exact_int64() -> None: # Python: 5-7 == -2, 3-1 == 2, (2**64-1)-(2**63) == 2**63-1 left = np.asarray([5, 3, U64_MAX], dtype=np.uint64) right = np.asarray([7, 1, 2**63], dtype=np.uint64) - out = _vec_arith(xp, int(ArithOp.SUB), left, right) + out = arith._vec_arith(xp, int(arith.ArithOp.SUB), left, right) assert out.dtype == np.dtype("int64") assert np.array_equal(out, [-2, 2, I64_MAX]) @@ -57,7 +55,7 @@ def test_vec_arith_uint64_sub_all_nonnegative_stays_uint64() -> None: # Python: 7-5 == 2, 3-1 == 2, 10-2 == 8 — all non-negative left = np.asarray([7, 3, 10], dtype=np.uint64) right = np.asarray([5, 1, 2], dtype=np.uint64) - out = _vec_arith(xp, int(ArithOp.SUB), left, right) + out = arith._vec_arith(xp, int(arith.ArithOp.SUB), left, right) assert out.dtype == np.dtype("uint64") assert np.array_equal(out, [2, 2, 8]) @@ -66,7 +64,7 @@ def test_vec_arith_uint64_mod_negative_literal_is_exact_int64() -> None: # Python: 10 % -3 == -2, 7 % -3 == -2, (2**64-1) % -3 == 0, 0 % -3 == 0 # (sign follows the divisor; 2**64-1 is divisible by 3) a = np.asarray([10, 7, U64_MAX, 0], dtype=np.uint64) - out = _vec_arith(xp, int(ArithOp.MOD), a, -3) + out = arith._vec_arith(xp, int(arith.ArithOp.MOD), a, -3) assert out.dtype == np.dtype("int64") assert np.array_equal(out, [-2, -2, 0, 0]) @@ -75,7 +73,7 @@ def test_vec_arith_uint64_floordiv_negative_literal_is_exact_int64() -> None: # Python: 10 // -3 == -4, 7 // -3 == -3, (2**64-1) // -3 == -6148914691236517205, # 0 // -3 == 0 (floor division rounds toward minus infinity) a = np.asarray([10, 7, U64_MAX, 0], dtype=np.uint64) - out = _vec_arith(xp, int(ArithOp.FLOORDIV), a, -3) + out = arith._vec_arith(xp, int(arith.ArithOp.FLOORDIV), a, -3) assert out.dtype == np.dtype("int64") assert np.array_equal(out, [-4, -3, -6148914691236517205, 0]) @@ -83,12 +81,14 @@ def test_vec_arith_uint64_floordiv_negative_literal_is_exact_int64() -> None: def test_vec_arith_div_is_always_float() -> None: # true division is float even for all-integer operands (strict # backends require floating-point operands) - out_i = _vec_arith(xp, int(ArithOp.DIV), np.asarray([1, 7, 9]), np.asarray([2, 2, 2])) + out_i = arith._vec_arith( + xp, int(arith.ArithOp.DIV), np.asarray([1, 7, 9]), np.asarray([2, 2, 2]) + ) assert out_i.dtype == np.dtype("float64") assert np.array_equal(out_i, [0.5, 3.5, 4.5]) - out_u = _vec_arith( + out_u = arith._vec_arith( xp, - int(ArithOp.DIV), + int(arith.ArithOp.DIV), np.asarray([9, 3], dtype=np.uint64), np.asarray([3, 2], dtype=np.uint64), ) @@ -98,7 +98,7 @@ def test_vec_arith_div_is_always_float() -> None: def test_vec_arith_neg_uint64_fits_int64() -> None: # Python: -5, -0 == 0 — magnitudes up to 2**63 fit int64 exactly - out = _vec_arith(xp, int(ArithOp.NEG), np.asarray([5, 0], dtype=np.uint64)) + out = arith._vec_arith(xp, int(arith.ArithOp.NEG), np.asarray([5, 0], dtype=np.uint64)) assert out.dtype == np.dtype("int64") assert np.array_equal(out, [-5, 0]) @@ -112,7 +112,7 @@ def test_vec_minmax_uint64_signed_min_fits_int64() -> None: # value, so they clamp to int64-max u = np.asarray([U64_MAX, 5, 100], dtype=np.uint64) s = np.asarray([-1, 3, -100], dtype=np.int64) - out = _vec_minmax(xp, True, u, s) + out = minmax._vec_minmax(xp, True, u, s) assert out.dtype == np.dtype("int64") assert np.array_equal(out, [-1, 3, -100]) @@ -123,7 +123,7 @@ def test_vec_minmax_uint64_signed_max_clamps_negatives() -> None: # so they clamp to 0 and the result stays uint64 u = np.asarray([U64_MAX, 5, 100], dtype=np.uint64) s = np.asarray([-1, 3, -100], dtype=np.int64) - out = _vec_minmax(xp, False, u, s) + out = minmax._vec_minmax(xp, False, u, s) assert out.dtype == np.dtype("uint64") assert np.array_equal(out, [U64_MAX, 5, 100]) @@ -133,10 +133,10 @@ def test_vec_minmax_bool_only_selection_returns_bool() -> None: # min(True, True) == True; max(True, True) == True a = np.asarray([True, False, True]) b = np.asarray([False, False, True]) - out = _vec_minmax(xp, True, a, b) + out = minmax._vec_minmax(xp, True, a, b) assert out.dtype == np.dtype("bool") assert np.array_equal(out, [False, False, True]) - out_lit = _vec_minmax(xp, False, np.asarray([True, False]), True) + out_lit = minmax._vec_minmax(xp, False, np.asarray([True, False]), True) assert out_lit.dtype == np.dtype("bool") assert np.array_equal(out_lit, [True, True]) @@ -146,14 +146,14 @@ def test_vec_minmax_bool_only_selection_returns_bool() -> None: def test_fits_dtype_edge_cases() -> None: # float32 cannot represent 2**24+1 (rounds to 16777216.0) - assert not _fits_dtype(16777217, np.dtype("float32")) - assert _fits_dtype(16777216, np.dtype("float32")) + assert not dtype._fits_dtype(16777217, np.dtype("float32")) + assert dtype._fits_dtype(16777216, np.dtype("float32")) # -1 is outside [0, 2**64) - assert not _fits_dtype(-1, np.dtype("uint64")) - assert _fits_dtype(-1, np.dtype("int64")) + assert not dtype._fits_dtype(-1, np.dtype("uint64")) + assert dtype._fits_dtype(-1, np.dtype("int64")) # only actual bools fit a bool dtype - assert _fits_dtype(True, np.dtype("bool")) - assert not _fits_dtype(2, np.dtype("bool")) + assert dtype._fits_dtype(True, np.dtype("bool")) + assert not dtype._fits_dtype(2, np.dtype("bool")) def test_fits_dtype_defensive_branches() -> None: @@ -164,17 +164,17 @@ def __init__(self, name: str) -> None: def __str__(self) -> str: return self._name - assert not _fits_dtype(1.5, _FakeDtype("float24")) # exotic width: promote - assert not _fits_dtype("not-a-number", np.dtype("int64")) - assert _fits_dtype(True, np.dtype("bool")) - assert not _fits_dtype(1, np.dtype("bool")) + assert not dtype._fits_dtype(1.5, _FakeDtype("float24")) # exotic width: promote + assert not dtype._fits_dtype("not-a-number", np.dtype("int64")) + assert dtype._fits_dtype(True, np.dtype("bool")) + assert not dtype._fits_dtype(1, np.dtype("bool")) def test_describe_dtype_classifies_uint64_as_float64_class() -> None: # uint64 fits no signed int, so it promotes as float64: (True, 64) - assert _describe_dtype(np.asarray([1], dtype=np.uint64)) == (True, 64) - assert _describe_dtype(1.5) == (True, 64) - assert _describe_dtype(np.asarray([1], dtype=np.int32)) == (False, 32) + assert dtype._describe_dtype(np.asarray([1], dtype=np.uint64)) == (True, 64) + assert dtype._describe_dtype(1.5) == (True, 64) + assert dtype._describe_dtype(np.asarray([1], dtype=np.int32)) == (False, 32) # small int literals count as the minimal dtype that fits - assert _describe_dtype(5) == (False, 8) - assert _describe_dtype(np.asarray([True])) == (False, 8) + assert dtype._describe_dtype(5) == (False, 8) + assert dtype._describe_dtype(np.asarray([True])) == (False, 8) diff --git a/tests/unit/test_validate.py b/tests/unit/test_validate.py index a9528b5..564eee9 100644 --- a/tests/unit/test_validate.py +++ b/tests/unit/test_validate.py @@ -8,28 +8,28 @@ from typing import Any import pytest -from support import make_module +import support -from array_vectorize import vectorize -from array_vectorize.errors import VectorizationError -from array_vectorize.frontend.extract import extract_function -from array_vectorize.frontend.validate import _EXPR_MESSAGES, validate +import array_vectorize +from array_vectorize import errors +from array_vectorize.frontend import extract +from array_vectorize.frontend import validate as validate_mod def make_fn(body: str, signature: str = "x", type_params: str = "") -> Callable[..., Any]: """Define a function with the given (indented) body in a real temp module.""" - mod = make_module( + mod = support.make_module( "import math\n\n\ndef subject" + type_params + "(" + signature + "):\n" + body + "\n" ) return mod.subject def check(src: str) -> None: - validate(extract_function(make_fn(f" return {src}"))) + validate_mod.validate(extract.extract_function(make_fn(f" return {src}"))) def check_body(body: str) -> None: - validate(extract_function(make_fn(body))) + validate_mod.validate(extract.extract_function(make_fn(body))) # ------------------------------------------------------------------ accepted @@ -128,7 +128,7 @@ def test_accepted_statements() -> None: ], ) def test_rejected_expressions(src: str, msg: str) -> None: - with pytest.raises(VectorizationError, match=msg): + with pytest.raises(errors.VectorizationError, match=msg): check(src) @@ -167,12 +167,12 @@ def test_rejected_statements(body: str, msg: str | None) -> None: if msg is None: check_body(body) return - with pytest.raises(VectorizationError, match=msg): + with pytest.raises(errors.VectorizationError, match=msg): check_body(body) def test_math_attribute_not_in_subset() -> None: - with pytest.raises(VectorizationError, match="supported math subset"): + with pytest.raises(errors.VectorizationError, match="supported math subset"): check("math.fsum([x])") @@ -182,8 +182,8 @@ def subject(x: float) -> float: x = x - 1 return x - with pytest.raises(VectorizationError) as exc: - validate(extract_function(subject)) + with pytest.raises(errors.VectorizationError) as exc: + validate_mod.validate(extract.extract_function(subject)) err = exc.value assert err.diagnostics, "expected at least one diagnostic" d = err.diagnostics[0] @@ -200,8 +200,8 @@ def subject(x: float) -> float: x = x - 1 return x, y - with pytest.raises(VectorizationError) as exc: - validate(extract_function(subject)) + with pytest.raises(errors.VectorizationError) as exc: + validate_mod.validate(extract.extract_function(subject)) messages = [d.message for d in exc.value.diagnostics] assert any("subscripts" in m for m in messages) assert any("while" in m for m in messages) @@ -209,28 +209,28 @@ def subject(x: float) -> float: def test_keyword_args_in_calls_rejected() -> None: - with pytest.raises(VectorizationError, match="keyword arguments in calls"): + with pytest.raises(errors.VectorizationError, match="keyword arguments in calls"): check("math.copysign(x, y=2.0)") def test_lambda_direct_is_accepted() -> None: - info = extract_function(lambda x: x + 1.0) # type: ignore[arg-type] - validate(info) + info = extract.extract_function(lambda x: x + 1.0) # type: ignore[arg-type] + validate_mod.validate(info) def test_docstring_only_first() -> None: check_body(' """doc"""\n return x') - with pytest.raises(VectorizationError, match="docstring"): + with pytest.raises(errors.VectorizationError, match="docstring"): check_body(' y = x\n """not first"""\n return y') def test_for_target_must_be_name() -> None: - with pytest.raises(VectorizationError, match="single name"): + with pytest.raises(errors.VectorizationError, match="single name"): check_body(" for i, j in range(3):\n pass\n return x") def test_nonlocal_rejected() -> None: - mod = make_module( + mod = support.make_module( "def outer():\n" " g = 1.0\n" " def subject(x):\n" @@ -238,8 +238,8 @@ def test_nonlocal_rejected() -> None: " return x\n" " return subject\n" ) - with pytest.raises(VectorizationError, match="nonlocal"): - validate(extract_function(mod.outer())) + with pytest.raises(errors.VectorizationError, match="nonlocal"): + validate_mod.validate(extract.extract_function(mod.outer())) # ------------------------------------------- Python 3.11-3.14 new syntax @@ -252,20 +252,20 @@ def test_nonlocal_rejected() -> None: def test_try_star_rejected() -> None: # except* is valid syntax on every supported interpreter (3.11+) - with pytest.raises(VectorizationError, match=r"except\*"): + with pytest.raises(errors.VectorizationError, match=r"except\*"): check_body(" try:\n y = x\n except* ValueError:\n y = 0\n return y") def test_type_alias_statement_rejected() -> None: # `type X = ...` (PEP 695) is valid syntax on every supported interpreter - with pytest.raises(VectorizationError, match="type alias statements"): + with pytest.raises(errors.VectorizationError, match="type alias statements"): check_body(" type Alias = int\n return x") def test_generic_function_rejected() -> None: # def subject[T](x) (PEP 695): type parameters on the extracted function - with pytest.raises(VectorizationError, match="type parameters"): - validate(extract_function(make_fn(" return x", type_params="[T]"))) + with pytest.raises(errors.VectorizationError, match="type parameters"): + validate_mod.validate(extract.extract_function(make_fn(" return x", type_params="[T]"))) def test_template_string_table_entry_is_version_gated() -> None: @@ -274,10 +274,10 @@ def test_template_string_table_entry_is_version_gated() -> None: template_node = getattr(ast, "TemplateStr", None) if sys.version_info >= (3, 14): assert template_node is not None - assert template_node in _EXPR_MESSAGES + assert template_node in validate_mod._EXPR_MESSAGES else: assert template_node is None - assert template_node not in _EXPR_MESSAGES + assert template_node not in validate_mod._EXPR_MESSAGES @pytest.mark.skipif(sys.version_info < (3, 14), reason="t-string syntax requires 3.14") @@ -285,7 +285,7 @@ def test_template_string_rejected() -> None: # t-strings (PEP 750) only exist in the 3.14 grammar; on 3.12/3.13 the # source itself would be a SyntaxError, so the rejection is only testable # here (the validator's guarded message-table entry covers both) - with pytest.raises(VectorizationError, match="t-strings"): + with pytest.raises(errors.VectorizationError, match="t-strings"): check('t"x"') @@ -309,38 +309,48 @@ def test_template_string_rejected() -> None: ) def test_validator_branches(body: str, msg: str | None) -> None: if msg is None: - validate(extract_function(make_fn(body, signature="x, y=2.0"))) + validate_mod.validate(extract.extract_function(make_fn(body, signature="x, y=2.0"))) return - with pytest.raises(VectorizationError, match=msg): - validate(extract_function(make_fn(body, signature="x, y=2.0"))) + with pytest.raises(errors.VectorizationError, match=msg): + validate_mod.validate(extract.extract_function(make_fn(body, signature="x, y=2.0"))) def test_attribute_assignment() -> None: - with pytest.raises(VectorizationError, match="attribute assignment"): - validate( - extract_function(make_fn(" a = x\n a.f = 1\n return a", signature="x, y=2.0")) + with pytest.raises(errors.VectorizationError, match="attribute assignment"): + validate_mod.validate( + extract.extract_function( + make_fn(" a = x\n a.f = 1\n return a", signature="x, y=2.0") + ) ) def test_generic_attribute_access_rejected() -> None: - with pytest.raises(VectorizationError, match="attribute access"): - validate(extract_function(make_fn(" return x.real", signature="x, y=2.0"))) + with pytest.raises(errors.VectorizationError, match="attribute access"): + validate_mod.validate( + extract.extract_function(make_fn(" return x.real", signature="x, y=2.0")) + ) def test_starred_call_argument_rejected() -> None: - with pytest.raises(VectorizationError, match="starred call arguments"): - validate(extract_function(make_fn(" return math.hypot(*[x, y])", signature="x, y=2.0"))) + with pytest.raises(errors.VectorizationError, match="starred call arguments"): + validate_mod.validate( + extract.extract_function( + make_fn(" return math.hypot(*[x, y])", signature="x, y=2.0") + ) + ) def test_unknown_expression_falls_through() -> None: # Starred in a non-call context hits the generic expression message - with pytest.raises(VectorizationError, match="not supported"): - validate(extract_function(make_fn(" return x + (y := 1)", signature="x, y=2.0"))) + with pytest.raises(errors.VectorizationError, match="not supported"): + validate_mod.validate( + extract.extract_function(make_fn(" return x + (y := 1)", signature="x, y=2.0")) + ) def test_del_is_rejected_by_vectorize() -> None: - with pytest.raises(VectorizationError, match="statements are not supported"): - vectorize(make_fn(" del x\n return 1.0", signature="x, y=2.0")) + with pytest.raises(errors.VectorizationError, match="statements are not supported"): + array_vectorize.vectorize(make_fn(" del x\n return 1.0", signature="x, y=2.0")) # --------------------------------------------------------- validator paths @@ -353,9 +363,11 @@ def test_del_is_rejected_by_vectorize() -> None: ], ) def test_validator_ok(body: str) -> None: - extract_function(make_fn(body)) + extract.extract_function(make_fn(body)) def test_statements_after_return_in_branch_rejected() -> None: - with pytest.raises(VectorizationError, match="after return"): - vectorize(make_fn(" if x > 0:\n return x\n y = 1\n return 0.0")) + with pytest.raises(errors.VectorizationError, match="after return"): + array_vectorize.vectorize( + make_fn(" if x > 0:\n return x\n y = 1\n return 0.0") + ) From 02ef4c802f3e891f03958d2f557234ab1b24e906 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 05:12:16 +0200 Subject: [PATCH 12/17] style: one import per line (google 3.13) --- pyproject.toml | 7 +++ src/array_vectorize/__init__.py | 3 +- src/array_vectorize/api.py | 8 ++-- src/array_vectorize/emit.py | 2 +- src/array_vectorize/fallback.py | 2 +- src/array_vectorize/frontend/__init__.py | 18 ++++---- src/array_vectorize/frontend/extract.py | 16 +++---- src/array_vectorize/frontend/info.py | 2 +- src/array_vectorize/frontend/lambda_id.py | 2 +- src/array_vectorize/frontend/validate.py | 3 +- src/array_vectorize/fuzz.py | 6 +-- src/array_vectorize/ir/__init__.py | 43 ++++++++++--------- src/array_vectorize/lower/__init__.py | 7 ++- src/array_vectorize/lower/env.py | 8 ++-- src/array_vectorize/lower/kinds.py | 2 +- src/array_vectorize/lower/types.py | 2 +- src/array_vectorize/optimize/cse.py | 3 +- src/array_vectorize/pipeline.py | 5 ++- src/array_vectorize/runtime/__init__.py | 3 +- src/array_vectorize/runtime/registry.py | 3 +- src/array_vectorize/verify.py | 5 ++- tests/backends/test_jax_torch_compilation.py | 2 +- tests/differential/test_scalar_equivalence.py | 2 +- tests/perf/test_bench.py | 3 +- tests/perf/test_speed_gate.py | 2 +- tests/support.py | 2 +- tests/unit/test_codegen.py | 3 +- tests/unit/test_docstring.py | 3 +- tests/unit/test_lower.py | 5 ++- tests/unit/test_optimize.py | 3 +- tests/unit/test_runtime.py | 4 +- tests/unit/test_validate.py | 2 +- 32 files changed, 105 insertions(+), 76 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2c9c1aa..195ba0a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -137,6 +137,13 @@ aliases = { numpy = "np" } # holds the NumPy-style docstring subject "tests/behavior/test_docstrings.py" = ["D4"] +[tool.ruff.lint.isort] +# Google §3.13 import formatting +force-single-line = true +single-line-exclusions = ["typing", "collections.abc"] +force-sort-within-sections = true +order-by-type = false + [tool.cleanporter] scope = "all" # golden cases are pinned compiler output; marimo owns the notebook format diff --git a/src/array_vectorize/__init__.py b/src/array_vectorize/__init__.py index 05bc573..16b22ca 100644 --- a/src/array_vectorize/__init__.py +++ b/src/array_vectorize/__init__.py @@ -6,7 +6,8 @@ from __future__ import annotations -from array_vectorize.api import get_source, vectorize +from array_vectorize.api import get_source +from array_vectorize.api import vectorize from array_vectorize.errors import VectorizationError __version__ = "0.1.0" diff --git a/src/array_vectorize/api.py b/src/array_vectorize/api.py index 3159bce..e1739d8 100644 --- a/src/array_vectorize/api.py +++ b/src/array_vectorize/api.py @@ -15,13 +15,15 @@ from __future__ import annotations -import types -import warnings from collections.abc import Callable +import types from typing import Any +import warnings -from array_vectorize import compat, errors, pipeline +from array_vectorize import compat +from array_vectorize import errors from array_vectorize import fallback as fallback_mod +from array_vectorize import pipeline from array_vectorize import verify as verify_mod from array_vectorize.frontend import extract diff --git a/src/array_vectorize/emit.py b/src/array_vectorize/emit.py index a0263b0..3a4966b 100644 --- a/src/array_vectorize/emit.py +++ b/src/array_vectorize/emit.py @@ -12,9 +12,9 @@ from __future__ import annotations +from collections.abc import Callable import inspect import linecache -from collections.abc import Callable from typing import Any, cast from array_vectorize.ir import ssa diff --git a/src/array_vectorize/fallback.py b/src/array_vectorize/fallback.py index 79087e9..aab7b7e 100644 --- a/src/array_vectorize/fallback.py +++ b/src/array_vectorize/fallback.py @@ -8,9 +8,9 @@ from __future__ import annotations +from collections.abc import Callable import functools import inspect -from collections.abc import Callable from typing import Any import array_api_compat diff --git a/src/array_vectorize/frontend/__init__.py b/src/array_vectorize/frontend/__init__.py index a00608e..2e1f5c1 100644 --- a/src/array_vectorize/frontend/__init__.py +++ b/src/array_vectorize/frontend/__init__.py @@ -8,16 +8,16 @@ from __future__ import annotations from array_vectorize.frontend.extract import extract_function -from array_vectorize.frontend.info import FunctionInfo, Param, ParamKind +from array_vectorize.frontend.info import FunctionInfo +from array_vectorize.frontend.info import Param +from array_vectorize.frontend.info import ParamKind from array_vectorize.frontend.lambda_id import _find_target -from array_vectorize.frontend.tables import ( - BUILTIN_CASTS, - BUILTIN_FOLDS, - BUILTIN_UNARY, - MATH_CONSTS, - MATH_FUNCS, - MATH_SPECIAL, -) +from array_vectorize.frontend.tables import BUILTIN_CASTS +from array_vectorize.frontend.tables import BUILTIN_FOLDS +from array_vectorize.frontend.tables import BUILTIN_UNARY +from array_vectorize.frontend.tables import MATH_CONSTS +from array_vectorize.frontend.tables import MATH_FUNCS +from array_vectorize.frontend.tables import MATH_SPECIAL __all__ = [ "BUILTIN_CASTS", diff --git a/src/array_vectorize/frontend/extract.py b/src/array_vectorize/frontend/extract.py index f9d6e45..2b5dc08 100644 --- a/src/array_vectorize/frontend/extract.py +++ b/src/array_vectorize/frontend/extract.py @@ -3,25 +3,25 @@ from __future__ import annotations import ast +from collections.abc import Callable import inspect import linecache import math import textwrap import types -from collections.abc import Callable from typing import Any -from array_vectorize import compat, errors +from array_vectorize import compat +from array_vectorize import errors from array_vectorize.frontend import info as info_mod -from array_vectorize.frontend import lambda_id, tables +from array_vectorize.frontend import lambda_id +from array_vectorize.frontend import tables # re-export for the module's public surface: __all__ names these, so the # from-import is what keeps them importable from frontend.extract -from array_vectorize.frontend.info import ( # cleanporter: ignore[CP001] - FunctionInfo, - Param, - ParamKind, -) +from array_vectorize.frontend.info import FunctionInfo # cleanporter: ignore[CP001] +from array_vectorize.frontend.info import Param # cleanporter: ignore[CP001] +from array_vectorize.frontend.info import ParamKind # cleanporter: ignore[CP001] __all__ = ["FunctionInfo", "Param", "ParamKind", "extract_function", "resolve_original"] diff --git a/src/array_vectorize/frontend/info.py b/src/array_vectorize/frontend/info.py index eae7325..14b14cd 100644 --- a/src/array_vectorize/frontend/info.py +++ b/src/array_vectorize/frontend/info.py @@ -3,8 +3,8 @@ from __future__ import annotations import ast -import dataclasses from collections.abc import Callable +import dataclasses from typing import Any, Literal __all__ = ["FunctionInfo", "Param", "ParamKind"] diff --git a/src/array_vectorize/frontend/lambda_id.py b/src/array_vectorize/frontend/lambda_id.py index f7b9270..3c77f9e 100644 --- a/src/array_vectorize/frontend/lambda_id.py +++ b/src/array_vectorize/frontend/lambda_id.py @@ -9,9 +9,9 @@ from __future__ import annotations import ast +from collections.abc import Iterator import math import types -from collections.abc import Iterator from typing import Any from array_vectorize import errors diff --git a/src/array_vectorize/frontend/validate.py b/src/array_vectorize/frontend/validate.py index 255d029..b0712bd 100644 --- a/src/array_vectorize/frontend/validate.py +++ b/src/array_vectorize/frontend/validate.py @@ -10,7 +10,8 @@ import ast from array_vectorize import errors -from array_vectorize.frontend import info, tables +from array_vectorize.frontend import info +from array_vectorize.frontend import tables __all__ = ["validate"] diff --git a/src/array_vectorize/fuzz.py b/src/array_vectorize/fuzz.py index 25ffcce..5be2802 100644 --- a/src/array_vectorize/fuzz.py +++ b/src/array_vectorize/fuzz.py @@ -23,10 +23,8 @@ # load-bearing re-export: tests/unit/test_fuzz.py monkeypatches # array_vectorize.fuzz.vectorize, so the name must stay bound here -from array_vectorize import ( - errors, - vectorize, # cleanporter: ignore[CP003] load-bearing -) +from array_vectorize import errors +from array_vectorize import vectorize # cleanporter: ignore[CP003] load-bearing _UNARY_MATH = [ f"math.{name}" diff --git a/src/array_vectorize/ir/__init__.py b/src/array_vectorize/ir/__init__.py index bd56763..f370a25 100644 --- a/src/array_vectorize/ir/__init__.py +++ b/src/array_vectorize/ir/__init__.py @@ -2,27 +2,28 @@ from __future__ import annotations -from array_vectorize.ir.nodes import ( - Binding, - BinOp, - Call, - Compare, - DType, - FuncCall, - Kind, - Literal, - Logical, - Loop, - Node, - Program, - Ref, - Stmt, - UnaryOp, - Where, - is_bool, -) -from array_vectorize.ir.ssa import RESERVED_NAMES, SSAEnv, generated_name -from array_vectorize.ir.walk import children, rewrite +from array_vectorize.ir.nodes import Binding +from array_vectorize.ir.nodes import BinOp +from array_vectorize.ir.nodes import Call +from array_vectorize.ir.nodes import Compare +from array_vectorize.ir.nodes import DType +from array_vectorize.ir.nodes import FuncCall +from array_vectorize.ir.nodes import is_bool +from array_vectorize.ir.nodes import Kind +from array_vectorize.ir.nodes import Literal +from array_vectorize.ir.nodes import Logical +from array_vectorize.ir.nodes import Loop +from array_vectorize.ir.nodes import Node +from array_vectorize.ir.nodes import Program +from array_vectorize.ir.nodes import Ref +from array_vectorize.ir.nodes import Stmt +from array_vectorize.ir.nodes import UnaryOp +from array_vectorize.ir.nodes import Where +from array_vectorize.ir.ssa import generated_name +from array_vectorize.ir.ssa import RESERVED_NAMES +from array_vectorize.ir.ssa import SSAEnv +from array_vectorize.ir.walk import children +from array_vectorize.ir.walk import rewrite __all__ = [ "RESERVED_NAMES", diff --git a/src/array_vectorize/lower/__init__.py b/src/array_vectorize/lower/__init__.py index c83460b..97e1e32 100644 --- a/src/array_vectorize/lower/__init__.py +++ b/src/array_vectorize/lower/__init__.py @@ -19,9 +19,12 @@ from array_vectorize.errors import VectorizationError from array_vectorize.frontend.info import FunctionInfo -from array_vectorize.ir import Node, Program, generated_name +from array_vectorize.ir import generated_name +from array_vectorize.ir import Node +from array_vectorize.ir import Program from array_vectorize.lower.helpers import _HelperLowerer -from array_vectorize.lower.types import HelperVectorizer, LoweredFunction +from array_vectorize.lower.types import HelperVectorizer +from array_vectorize.lower.types import LoweredFunction from array_vectorize.runtime.registry import RUNTIME_HELPERS __all__ = ["HelperVectorizer", "LoweredFunction", "lower_function"] diff --git a/src/array_vectorize/lower/env.py b/src/array_vectorize/lower/env.py index 1c05e3c..9017069 100644 --- a/src/array_vectorize/lower/env.py +++ b/src/array_vectorize/lower/env.py @@ -12,11 +12,13 @@ import ast from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import Any, TYPE_CHECKING -from array_vectorize import errors, ir +from array_vectorize import errors +from array_vectorize import ir from array_vectorize.frontend import info -from array_vectorize.lower import kinds, types +from array_vectorize.lower import kinds +from array_vectorize.lower import types from array_vectorize.runtime import registry __all__ = ["_LowererBase"] diff --git a/src/array_vectorize/lower/kinds.py b/src/array_vectorize/lower/kinds.py index ed2daac..7214a69 100644 --- a/src/array_vectorize/lower/kinds.py +++ b/src/array_vectorize/lower/kinds.py @@ -42,8 +42,8 @@ from __future__ import annotations -import dataclasses from collections.abc import Mapping +import dataclasses from array_vectorize import ir diff --git a/src/array_vectorize/lower/types.py b/src/array_vectorize/lower/types.py index c874027..4edd7d8 100644 --- a/src/array_vectorize/lower/types.py +++ b/src/array_vectorize/lower/types.py @@ -8,8 +8,8 @@ from __future__ import annotations -import dataclasses from collections.abc import Callable +import dataclasses from typing import Any from array_vectorize import ir diff --git a/src/array_vectorize/optimize/cse.py b/src/array_vectorize/optimize/cse.py index 0700b1e..4cea993 100644 --- a/src/array_vectorize/optimize/cse.py +++ b/src/array_vectorize/optimize/cse.py @@ -5,7 +5,8 @@ import collections from array_vectorize import ir -from array_vectorize.ir import ssa, walk +from array_vectorize.ir import ssa +from array_vectorize.ir import walk __all__ = ["cse"] diff --git a/src/array_vectorize/pipeline.py b/src/array_vectorize/pipeline.py index 3a3df14..ddfe1c1 100644 --- a/src/array_vectorize/pipeline.py +++ b/src/array_vectorize/pipeline.py @@ -12,8 +12,11 @@ from collections.abc import Callable from typing import Any -from array_vectorize import codegen, emit, lower, verify +from array_vectorize import codegen +from array_vectorize import emit +from array_vectorize import lower from array_vectorize import optimize as optimize_mod +from array_vectorize import verify from array_vectorize.frontend import extract from array_vectorize.frontend import validate as validate_mod diff --git a/src/array_vectorize/runtime/__init__.py b/src/array_vectorize/runtime/__init__.py index 86ce519..16bd655 100644 --- a/src/array_vectorize/runtime/__init__.py +++ b/src/array_vectorize/runtime/__init__.py @@ -7,7 +7,8 @@ from __future__ import annotations -from array_vectorize.runtime.arith import ArithOp, _vec_arith +from array_vectorize.runtime.arith import _vec_arith +from array_vectorize.runtime.arith import ArithOp from array_vectorize.runtime.dtype import _fits_dtype from array_vectorize.runtime.minmax import _vec_minmax from array_vectorize.runtime.registry import RUNTIME_HELPERS diff --git a/src/array_vectorize/runtime/registry.py b/src/array_vectorize/runtime/registry.py index f443c89..f639a24 100644 --- a/src/array_vectorize/runtime/registry.py +++ b/src/array_vectorize/runtime/registry.py @@ -11,7 +11,8 @@ from collections.abc import Callable from typing import Any -from array_vectorize.runtime import arith, minmax +from array_vectorize.runtime import arith +from array_vectorize.runtime import minmax # lowering goes through the registry seam (docs/architecture.md): # lower/expressions.py imports ArithOp from here, not from runtime.arith diff --git a/src/array_vectorize/verify.py b/src/array_vectorize/verify.py index 7235f71..08e1e64 100644 --- a/src/array_vectorize/verify.py +++ b/src/array_vectorize/verify.py @@ -8,13 +8,14 @@ from __future__ import annotations -import warnings from collections.abc import Callable from typing import Any +import warnings import array_api_compat -from array_vectorize import compat, errors +from array_vectorize import compat +from array_vectorize import errors __all__ = ["verify_match"] diff --git a/tests/backends/test_jax_torch_compilation.py b/tests/backends/test_jax_torch_compilation.py index 51743e6..d8eccfb 100644 --- a/tests/backends/test_jax_torch_compilation.py +++ b/tests/backends/test_jax_torch_compilation.py @@ -21,9 +21,9 @@ from __future__ import annotations +from collections.abc import Callable, Sequence import importlib.util import pathlib -from collections.abc import Callable, Sequence from typing import Any import numpy as np diff --git a/tests/differential/test_scalar_equivalence.py b/tests/differential/test_scalar_equivalence.py index fdf5647..901e04e 100644 --- a/tests/differential/test_scalar_equivalence.py +++ b/tests/differential/test_scalar_equivalence.py @@ -6,8 +6,8 @@ import corpus import hypothesis -import numpy as np from hypothesis import strategies as st +import numpy as np import array_vectorize diff --git a/tests/perf/test_bench.py b/tests/perf/test_bench.py index ff829a4..c2ba827 100644 --- a/tests/perf/test_bench.py +++ b/tests/perf/test_bench.py @@ -24,7 +24,8 @@ import pytest import array_vectorize -from array_vectorize.runtime import arith, minmax +from array_vectorize.runtime import arith +from array_vectorize.runtime import minmax pytestmark = pytest.mark.slow diff --git a/tests/perf/test_speed_gate.py b/tests/perf/test_speed_gate.py index 133ef8a..c4da1ae 100644 --- a/tests/perf/test_speed_gate.py +++ b/tests/perf/test_speed_gate.py @@ -14,9 +14,9 @@ from __future__ import annotations +from collections.abc import Callable import math import time -from collections.abc import Callable import numpy as np import pytest diff --git a/tests/support.py b/tests/support.py index 0156d64..e37783b 100644 --- a/tests/support.py +++ b/tests/support.py @@ -10,11 +10,11 @@ from __future__ import annotations +from collections.abc import Callable import importlib.util import itertools import pathlib import tempfile -from collections.abc import Callable from typing import Any import corpus diff --git a/tests/unit/test_codegen.py b/tests/unit/test_codegen.py index 0a5db38..295dcc1 100644 --- a/tests/unit/test_codegen.py +++ b/tests/unit/test_codegen.py @@ -4,7 +4,8 @@ import ast -from array_vectorize import codegen, ir +from array_vectorize import codegen +from array_vectorize import ir from array_vectorize.frontend import info from array_vectorize.lower import types diff --git a/tests/unit/test_docstring.py b/tests/unit/test_docstring.py index b7b4c6a..7ff8219 100644 --- a/tests/unit/test_docstring.py +++ b/tests/unit/test_docstring.py @@ -6,7 +6,8 @@ import pytest -from array_vectorize.codegen import docstring, emitter +from array_vectorize.codegen import docstring +from array_vectorize.codegen import emitter SOURCE = ( "def clamp(x, lo=0.0, hi=1.0):\n" diff --git a/tests/unit/test_lower.py b/tests/unit/test_lower.py index ba7bca7..cb37e8f 100644 --- a/tests/unit/test_lower.py +++ b/tests/unit/test_lower.py @@ -2,8 +2,8 @@ from __future__ import annotations -import math from collections.abc import Callable +import math from typing import Any import numpy as np @@ -11,7 +11,8 @@ import support import array_vectorize -from array_vectorize import errors, ir +from array_vectorize import errors +from array_vectorize import ir from array_vectorize import lower as lower_mod from array_vectorize.frontend import extract diff --git a/tests/unit/test_optimize.py b/tests/unit/test_optimize.py index 4839b07..87d16ce 100644 --- a/tests/unit/test_optimize.py +++ b/tests/unit/test_optimize.py @@ -11,7 +11,8 @@ # optimize's documented submodule/function name collision (the NOTE in # optimize/__init__.py): the from-import is the supported spelling -from array_vectorize.optimize import cse, dce # cleanporter: ignore[CP002] name collision +from array_vectorize.optimize import cse # cleanporter: ignore[CP002] name collision +from array_vectorize.optimize import dce # cleanporter: ignore[CP002] name collision from array_vectorize.optimize import dce as dce_pass # cleanporter: ignore[CP002] name collision diff --git a/tests/unit/test_runtime.py b/tests/unit/test_runtime.py index c90c79a..faba071 100644 --- a/tests/unit/test_runtime.py +++ b/tests/unit/test_runtime.py @@ -10,7 +10,9 @@ import numpy as np -from array_vectorize.runtime import arith, dtype, minmax +from array_vectorize.runtime import arith +from array_vectorize.runtime import dtype +from array_vectorize.runtime import minmax # the array namespace passed to the runtime helpers, as generated code # binds it diff --git a/tests/unit/test_validate.py b/tests/unit/test_validate.py index 564eee9..fcbabf4 100644 --- a/tests/unit/test_validate.py +++ b/tests/unit/test_validate.py @@ -3,8 +3,8 @@ from __future__ import annotations import ast -import sys from collections.abc import Callable +import sys from typing import Any import pytest From 3766c3ebf8340a480fc1534b59849d5f9b77dd95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 05:21:29 +0200 Subject: [PATCH 13/17] build: gate module-only imports in make lint --- AGENTS.md | 4 +++- Makefile | 3 +++ docs/development.md | 18 ++++++++++++++++-- pyproject.toml | 2 ++ 4 files changed, 24 insertions(+), 3 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 2a22898..efadcaf 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -46,7 +46,9 @@ Invariants: change, then review `git diff tests/golden/cases` line by line. The code follows the Google Python Style Guide, single deviation: -100-column lines. Ruff enforces it via `make lint` / pre-commit. +100-column lines. Imports follow Google §2.2 (import modules, not +names), enforced by cleanporter in `make lint` and `make fmt`; ruff +enforces the rest via `make lint` / pre-commit. ## Tests diff --git a/Makefile b/Makefile index 20ddd68..432ada1 100644 --- a/Makefile +++ b/Makefile @@ -33,6 +33,9 @@ fmt: lint: $(RUN) ruff check $(PY_PATHS) $(RUN) ruff format --check $(PY_PATHS) + # the import gate: exits 0 only when clean -- treat_unresolved_as_error + # makes unclassifiable imports fail too + $(RUN) cleanporter $(PY_PATHS) type: $(RUN) mypy diff --git a/docs/development.md b/docs/development.md index 3d9b132..9830365 100644 --- a/docs/development.md +++ b/docs/development.md @@ -22,11 +22,11 @@ dependency at its floor (`--resolution lowest-direct`) on Python 3.12. ## Commands ```bash -make check # ruff + mypy (strict) + pytest + coverage (gate: 95%) +make check # cleanporter + ruff + mypy (strict) + pytest + coverage (gate: 95%) make check PY=3.13 # any target on another Python (env: .venv-3.13) make lint type test # the same gates without coverage make lowest # the test suite at the minimum dependency versions -make fmt # format + autofix +make fmt # cleanporter --fix, then ruff format / check --fix / format make smoke # run scripts/smoke.py, examples/demo.py, the notebooks make fuzz # grammar fuzzer on the CI seeds make bench # pytest-benchmark suite (timings + dispatch gates) @@ -43,6 +43,10 @@ The default `make check` also runs the slow-marked performance gate; plain `pytest` skips nothing but keeps benchmarks disabled (`--benchmark-disable` is in `addopts`). +`make lint` includes the cleanporter import gate, and `make fmt` runs the +same check with `--fix`: cleanporter rewrites first, so ruff formats the +already-wrapped lines (its module prefixes push lines past 100 columns). + ## Testing layers | Layer | Where | @@ -74,6 +78,16 @@ conventions (`numpy` as `np`), absolute imports only (TID252, `ban-relative-imports = "all"`), and additionally BLE001, PLW0603, G004, PGH003 and PGH004 (evaluated: zero violations). +Imports follow Google §2.2 (import modules, not names), enforced by +[cleanporter](https://pypi.org/project/cleanporter/) in `make lint` and +`make fmt` (scope `all`; the golden cases and the notebooks are excluded +as generated and marimo-owned, `tests/corpus.py` is skipped as golden +test data, and `treat_unresolved_as_error` is on, so an import +cleanporter cannot classify fails the gate). The deliberate exceptions +are inline `# cleanporter: ignore[...]` suppressions, each naming its +finding code and carrying its reason on the import's line. One import +per line per §3.13 (ruff isort `force-single-line`). + Evaluated and deliberately skipped: TD (conflicts with Google's current TODO format) and PLC0415 (import-inside-function is a deliberate pattern in this compiler's test subjects). diff --git a/pyproject.toml b/pyproject.toml index 195ba0a..43f1bd3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -148,6 +148,8 @@ order-by-type = false scope = "all" # golden cases are pinned compiler output; marimo owns the notebook format exclude = ["tests/golden/cases/", "examples/notebooks/"] +# unclassifiable imports (CP002) fail the gate, not only plain violations +treat_unresolved_as_error = true [[tool.cleanporter.skip]] file = 'tests/corpus\.py' From a2cbb2e42f0c7187ccdd18065ff0c7d83dff1220 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 08:24:23 +0200 Subject: [PATCH 14/17] build: pin cleanporter to the upstream 0.6.0 release --- pyproject.toml | 5 ++--- uv.lock | 4 ++-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 43f1bd3..47d4644 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,9 +72,8 @@ explicit = true [tool.uv.sources] torch = { index = "pytorch-cpu" } -# local stand-in pin for the PEP 695 type-alias fix; swap to the GitHub -# URL + 0.6.0 tag when the upstream release is cut -cleanporter = { git = "file:///tmp/opencode/cleanporter", rev = "e53776c" } +# 0.6.0 is the first release carrying the PEP 695 type-alias fix +cleanporter = { git = "https://github.com/grAItools/cleanporter", tag = "v0.6.0" } [tool.hatch.build.targets.wheel] packages = ["src/array_vectorize"] diff --git a/uv.lock b/uv.lock index 4cc61d2..558cf58 100644 --- a/uv.lock +++ b/uv.lock @@ -92,7 +92,7 @@ backends = [ ] dev = [ { name = "array-api-strict", specifier = ">=2.3" }, - { name = "cleanporter", git = "file:///tmp/opencode/cleanporter?rev=e53776c" }, + { name = "cleanporter", git = "https://github.com/grAItools/cleanporter?tag=v0.6.0" }, { name = "coverage", extras = ["toml"], specifier = ">=7.6" }, { name = "hypothesis", specifier = ">=6.100" }, { name = "mypy", specifier = ">=1.20" }, @@ -271,7 +271,7 @@ wheels = [ [[package]] name = "cleanporter" version = "0.6.0" -source = { git = "file:///tmp/opencode/cleanporter?rev=e53776c#e53776c0898cf3fe0ac141fa2ce3ff4655213520" } +source = { git = "https://github.com/grAItools/cleanporter?tag=v0.6.0#61d18ec658e7e4a200ea18621d5734b6938c4c4c" } dependencies = [ { name = "libcst" }, ] From 729993a52c9548b368af1bac0c156a6a53f338bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 09:52:01 +0200 Subject: [PATCH 15/17] build: fail make fmt on cleanporter operational errors --- Makefile | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index 432ada1..2e5d073 100644 --- a/Makefile +++ b/Makefile @@ -19,13 +19,14 @@ install: $(RUN) pre-commit install # fixers run to completion whatever findings remain; `make lint` stays -# the gate. `-` on cleanporter: it exits 1 whenever anything remains it -# will not rewrite (no exit-zero flag). ruff check --fix gets -# --exit-zero: ruff cannot fix E501 itself, and cleanporter's module -# prefixes push lines past 100 columns; ruff format runs first and last -# so the fixes land on wrapped code +# the gate. cleanporter gets `test $$? -le 1`: it exits 1 whenever +# anything remains it will not rewrite (no exit-zero flag), but exit 2 +# is an operational error and must fail the recipe. ruff check --fix +# gets --exit-zero: ruff cannot fix E501 itself, and cleanporter's +# module prefixes push lines past 100 columns; ruff format runs first +# and last so the fixes land on wrapped code fmt: - -$(RUN) cleanporter --fix $(PY_PATHS) + $(RUN) cleanporter --fix $(PY_PATHS); test $$? -le 1 $(RUN) ruff format $(PY_PATHS) $(RUN) ruff check --fix --exit-zero $(PY_PATHS) $(RUN) ruff format $(PY_PATHS) From 7ffa6c38f681cf096a769de50d32cf2160552cf6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 09:52:02 +0200 Subject: [PATCH 16/17] docs: fix the cleanporter link and suppression-reason wording --- docs/development.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/development.md b/docs/development.md index 9830365..7222a01 100644 --- a/docs/development.md +++ b/docs/development.md @@ -79,13 +79,14 @@ conventions (`numpy` as `np`), absolute imports only (TID252, PGH003 and PGH004 (evaluated: zero violations). Imports follow Google §2.2 (import modules, not names), enforced by -[cleanporter](https://pypi.org/project/cleanporter/) in `make lint` and +[cleanporter](https://github.com/grAItools/cleanporter) in `make lint` and `make fmt` (scope `all`; the golden cases and the notebooks are excluded as generated and marimo-owned, `tests/corpus.py` is skipped as golden test data, and `treat_unresolved_as_error` is on, so an import cleanporter cannot classify fails the gate). The deliberate exceptions are inline `# cleanporter: ignore[...]` suppressions, each naming its -finding code and carrying its reason on the import's line. One import +finding code and carrying its reason in the suppression comment or the +comment block above it. One import per line per §3.13 (ruff isort `force-single-line`). Evaluated and deliberately skipped: TD (conflicts with Google's current From 4e78c5af52ebeb069a29a7cc2354f1a19cd44cac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enrique=20Gonz=C3=A1lez=20Paredes?= Date: Wed, 7 Oct 2026 09:52:04 +0200 Subject: [PATCH 17/17] ci: check PR commit subjects against Conventional Commits --- .github/workflows/ci.yml | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9ba5cb3..99c4fff 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -83,3 +83,30 @@ jobs: run: uv sync --locked --group docs - name: Build docs run: make docs + + # PR commit subjects follow Conventional Commits, the same rule the + # local commit-msg hook enforces (compilerla/conventional-pre-commit, + # pinned to the same tag); the full history is needed to walk the + # PR's commit range + commits: + if: github.event_name == 'pull_request' + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v7 + with: + fetch-depth: 0 + - uses: astral-sh/setup-uv@v10.2.0 + - name: Check commit subjects + run: | + base="origin/${{ github.base_ref }}" + status=0 + for sha in $(git rev-list --reverse --no-merges "$base"..HEAD); do + git log -1 --format=%s "$sha" > /tmp/commit-subject.txt + if ! uvx --from git+https://github.com/compilerla/conventional-pre-commit@v4.4.0 \ + conventional-pre-commit feat fix docs style refactor perf test build ci chore revert \ + /tmp/commit-subject.txt; then + echo "::error::$(git log -1 --oneline "$sha") is not a Conventional Commit" + status=1 + fi + done + exit $status