Merge branch 'main' into u/vineetgarg/block_activations

Author: vineet-g

.github/workflows/release.yml

@@ -52,14 +52,20 @@ jobs:
         # The published version comes from src/coreai_opt/_about.py, not the tag.
         # Fail early if they disagree so we never publish a mismatched/duplicate
         # version (PyPI uploads are immutable and cannot be overwritten).
+        # `print_version.py --release` computes the version exactly as the
+        # `make build` step below does, so this guard can't drift from what
+        # actually gets published.
         # Skipped on manual dry runs, where the ref is a branch, not a vX.Y.Z tag.
+        # Run via uv (installed above) so the interpreter satisfies
+        # requires-python whatever the runner image ships; see the Makefile's
+        # `version` target.
         if: github.event_name == 'push'
         run: |
           tag="${GITHUB_REF_NAME}"
-          version="$(python3 -c "import runpy; print(runpy.run_path('src/coreai_opt/_about.py')['__version__'])")"
+          version="$(uv run --no-config --no-project --python '>=3.11' scripts/make/print_version.py --release)"
           echo "tag=${tag}  package version=${version}"
           if [ "${tag}" != "v${version}" ]; then
-            echo "::error::Tag ${tag} does not match package version v${version} (src/coreai_opt/_about.py). Bump __version__ to match the release tag."
+            echo "::error::Tag ${tag} does not match package version v${version} (src/coreai_opt/_about.py). Update latest_released_version so the release it implies matches the tag."
             exit 1
           fi
       - name: Build wheel and sdist

.pre-commit-config.yaml

@@ -319,6 +319,20 @@ repos:
         types: [python]
         exclude: ^tests/
 
+  - repo: local
+    hooks:
+      - id: check-about-version
+        name: Check _about.py version fields
+        description: |
+          Check that _about.py's latest_released_version matches the repo's
+          latest release tag, and that __version__ is its last number plus
+          one, plus .dev0. Catches a release candidate that looks like a
+          release has already shipped when it hasn't.
+        entry: python scripts/pre_commit/check_about_version.py
+        language: system
+        files: (^|/)_about\.py$
+        pass_filenames: false
+
   - repo: local
     hooks:
       - id: towncrier-check

Makefile

@@ -95,6 +95,25 @@ SHELL_RC ?=
 # internal-only venv defaults.
 ENV_ALL_EXTRAS ?= true
 
+# Extra `uv pip install` arguments, applied to a venv after it is synced and to
+# each nox session venv (see ci/nox/noxfile.py). Use it to swap a dependency
+# version for one run without editing pyproject.toml or the lockfile — for
+# example a scheduled job testing against newer upstream builds. Empty (the
+# default) changes nothing.
+POST_INSTALL_PIP_ARGS ?=
+export POST_INSTALL_PIP_ARGS
+
+# Re-apply POST_INSTALL_PIP_ARGS to the venv at $(1).
+#
+# setup_env.sh already does this, so only use it when a recipe installs more
+# packages afterwards that could pull the old version back in (see env-all).
+# Expands to `true` when unset. The empty check uses `$(if ...)` instead of a
+# shell `[ -n "..." ]` test because the args contain their own quotes.
+# Usage: $(call post_install_pip,VENV_PATH)
+define post_install_pip
+$(if $(POST_INSTALL_PIP_ARGS),echo "Applying POST_INSTALL_PIP_ARGS to $(1)" && source $(1)/bin/activate && uv pip install $(POST_INSTALL_PIP_ARGS),true)
+endef
+
 # Local wheelhouse for pre-release wheels not yet on an index.
 # Exported so every recipe-level uv invocation (uv lock, uv sync, uv venv)
 # resolves matching packages from disk without an index lookup. The wildcard
@@ -212,23 +231,29 @@ env-all: _maybe_patch_pyproject
 	@$(SETUP_ENV) --venv $(VENV) --python-version $(PYTHON_VERSION) --all-groups
 	@$(call write_active_venv,$(VENV))
 	@$(ENV_ALL_EXTRAS)
+	@$(call post_install_pip,$(VENV))
 
 # =============================================================================
 # Build
 # =============================================================================
 
-# Build the canonical, publishable distribution (wheel + sdist) via the uv build
-# frontend. `--no-sources` ignores [tool.uv.sources], so the artifact doesn't
-# depend on uv-specific index overrides — the recommended way to build for
-# publication. This is what the release workflow runs.
+# Build the canonical, publishable distribution (wheel + sdist): the on-tree
+# version with any `.dev` suffix stripped (e.g. 0.2.2.dev0 -> 0.2.2), via
+# `uv build --no-sources`. `--no-sources` ignores [tool.uv.sources], so the
+# artifact doesn't depend on uv-specific index overrides — the recommended way
+# to build for publication. This is what the release workflow runs. Routed
+# through build.py (like build-dev) so both targets share one code path; set
+# COREAI_OPT_VERSION_EXTENSION to insert an extra release segment (see
+# RELEASE.md).
 build:
-	@uv build --no-sources
+	@$(call use_env,VENV) && uv run --no-sync --active python $(SCRIPTS)/make/build.py --no-sources
 
-# Build the development distribution with build.py (standard version; build.py
-# also supports a PEP 440 .dev version via --dev). Used by contributors and the
-# smoke tests.
+# Build a development distribution with build.py: the release base with a
+# unique, timestamped PEP 440 dev suffix (e.g. 0.2.2.dev202607231430+abc1234).
+# Used by contributors, the smoke tests, and the nightly pipeline. Set
+# DEV_VERSION=... to use an exact version instead.
 build-dev:
-	@$(call use_env,VENV) && uv run --no-sync --active python $(SCRIPTS)/make/build.py
+	@$(call use_env,VENV) && uv run --no-sync --active python $(SCRIPTS)/make/build.py --dev
 
 # =============================================================================
 # Code Quality
@@ -277,19 +302,24 @@ test-smoke:
 	uv run --no-sync --active nox -f $(MAKEFILE_DIR)ci/nox/noxfile.py -s smoke_tests -- $(PYTEST_ARGS) && \
 	echo "All smoke tests passed!"
 
-# Run tests on lowest supported PyTorch version (pass PYTEST_ARGS for custom flags)
-test-lowest-pytorch: env-lowest-torch
+# Run tests on lowest supported PyTorch version (pass PYTEST_ARGS for custom flags).
+# TORCH_GROUP is already exported, so setting it per target is enough for
+# use_env to pick the right torch build. Use `=`, not `:=`: an including
+# Makefile may change HIGHEST_TORCH_GROUP after this file is read.
+test-lowest-pytorch: TORCH_GROUP = $(LOWEST_TORCH_GROUP)
+test-lowest-pytorch:
 	@echo "Running tests on lowest PyTorch version supported..."
-	@source $(VENV_LOWEST_TORCH)/bin/activate && \
+	@$(call use_env,VENV_LOWEST_TORCH) && \
 	echo "Testing with lowest supported PyTorch versions" && \
 	uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
 	$(RUN_TESTS) $(PYTEST_ARGS) && \
 	echo "All tests passed!"
 
 # Run tests on highest supported PyTorch version (pass PYTEST_ARGS for custom flags)
-test-highest-pytorch: env-highest-torch
+test-highest-pytorch: TORCH_GROUP = $(HIGHEST_TORCH_GROUP)
+test-highest-pytorch:
 	@echo "Running tests on highest PyTorch version supported..."
-	@source $(VENV_HIGHEST_TORCH)/bin/activate && \
+	@$(call use_env,VENV_HIGHEST_TORCH) && \
 	echo "Testing with latest supported PyTorch versions" && \
 	uv run --no-sync --active python $(SCRIPTS)/make/log_versions.py && \
 	$(RUN_TESTS) $(PYTEST_ARGS) && \
@@ -326,9 +356,13 @@ distclean-all:
 set-auto-venv:
 	@$(SCRIPTS)/make/set_auto_venv.sh $(DEFAULT_VENV) $(SHELL_RC)
 
-# Show current version
+# Show the development version carried on the tree (e.g. 0.2.2.dev0), including
+# any COREAI_OPT_VERSION_EXTENSION (e.g. 0.2.2.1.dev0). Reads _about.py as plain
+# text, so no venv is needed — but `uv run --no-project` is still what guarantees
+# a >= 3.11 interpreter (a bare `python3` is 3.9 on stock macOS, and `python` may
+# not exist at all) without requiring `make env` first.
 version:
-	@python -c "exec(open('$(MAKEFILE_DIR)src/coreai_opt/_about.py').read()); print(__version__)"
+	@uv run --no-config --no-project --python '>=3.11' $(SCRIPTS)/make/print_version.py
 
 # =============================================================================
 # Documentation

RELEASE.md

@@ -1,13 +1,37 @@
 # Package Release Guide
 
-The OSS release process for Core AI Optimization is being defined. This page will document the workflow for publishing to PyPI once the public release infrastructure is finalized.
+The OSS release process for Core AI Optimization is still being defined. This page will document the workflow for publishing to PyPI once the public release infrastructure is finalized.
 
-Available locally:
+The following commands are available locally:
 
 ```bash
-make build      # build the package wheel
-make version    # show current version
+make build      # build the canonical, publishable wheel + sdist (uv build --no-sources)
+make build-dev  # build a timestamped dev wheel (e.g. 0.2.2.dev202607231430+abc1234)
+make version    # show the development version carried on the tree (e.g. 0.2.2.dev0)
 make clean      # remove build artifacts
 ```
 
+## Version scheme
+
+`main` always carries the version planned for the _next_ release. This ensures that ongoing development is never mistaken for an already-published version, and that a release can be stabilized, tested, and published on its own branch, independently of later changes on `main`. (The release-branch workflow itself — branch naming, tagging, and backporting fixes to `main` — will be documented separately in the release schedule doc; this section covers only the version-string mechanics.)
+
+`src/coreai_opt/_about.py` stores `latest_released_version` (the last tagged release, e.g. `"0.2.1"`) and computes `__version__` from it by incrementing its last number by one and adding a `.dev0` suffix (e.g. `"0.2.2.dev0"`). A pre-commit hook (`check-about-version`) verifies that `__version__` always follows this rule and that `latest_released_version` matches the repo's latest release tag. As a result, `__version__` can never look as though a release has shipped when it hasn't. The `.dev0` suffix is only a marker on the tree; it never appears in a built wheel.
+
+- `make build` builds the release that `__version__` implies, e.g. `0.2.2`. A release is cut by tagging it (`v0.2.2`); `latest_released_version` is then hard-coded to `"0.2.2"`, which bumps `__version__` to the next candidate (`0.2.3.dev0`).
+- `make build-dev` builds that same release but with a unique `.dev<UTC-timestamp>+<short-sha>` suffix instead. It is used by contributors, smoke tests, and the nightly pipeline. `DEV_VERSION=<version>` uses that version exactly instead.
+
+Sorting is preserved: `0.2.2.dev0 < 0.2.2.dev202607231430+abc1234 < 0.2.2`.
+
+### Extending the scheme downstream
+
+A repo that uses this one as a submodule and includes this `Makefile` — building one combined wheel from both trees — can add its own 4th number. Set `COREAI_OPT_VERSION_EXTENSION` to the number it's about to release next (e.g. `"1"` for its first release off a given OSS release, then `"2"` for the one after that). Then call `make build`, `make build-dev`, or `make version` unchanged:
+
+- `latest_released_version` `"0.2.1"` + extension `"1"` -> candidate: `0.2.1.1.dev0`
+- `make build` -> `0.2.1.1`
+- `make build-dev` -> `0.2.1.1.dev<UTC-timestamp>+<short-sha>`
+
+The extra number is used exactly as given (`scripts/release/release_utils.apply_version_extension`); `latest_released_version`'s own last number is only bumped for OSS's own `main`, when no extension is set.
+
+There is still only one `_about.py` (this package's own); the extra number is a plain string handled entirely in `scripts/release/release_utils.next_release_base` — no other file or package is involved.
+
 <!-- TODO: Document the chosen OSS release workflow (PyPI trusted publishing, twine upload, or uv publish). -->

changelog.d/52.fixed

@@ -0,0 +1 @@
+Fix per-channel activation quantization crashing with a shape-mismatch `RuntimeError` on MaxPool/AvgPool/AdaptiveAvgPool layers whose shared observer spans an axis the pool shrinks (e.g. a spatial axis under a stride>1 pool). Axes pooling never touches (batch, channel) keep working as per-channel; only the specific unsafe axis falls back to per-tensor, with a warning explaining why and how to pick a safe axis instead.

changelog.d/53.changed

@@ -0,0 +1 @@
+Carry the next planned release with a `.dev0` suffix on `main` (e.g. `0.2.2.dev0`). `make build` strips the suffix for a clean release and `make build-dev` builds a unique, timestamped dev wheel (`0.2.2.dev<timestamp>+<shortsha>`). A repo that vendors this one can insert one extra release segment via the `COREAI_OPT_VERSION_EXTENSION` environment variable to extend the version scheme.

ci/nox/noxfile.py

@@ -10,6 +10,7 @@
 """
 
 import os
+import shlex
 import sys
 from pathlib import Path
 
@@ -42,6 +43,12 @@
 # resolved against the working directory (the project root, set just below).
 SMOKE_TEST_DIST = os.environ.get("SMOKE_TEST_DIST")
 
+# Extra `uv pip install` arguments, applied to the session venv after the
+# package under test is installed. Use it to swap a dependency version for one
+# run without editing pyproject.toml or the lockfile. Split with shlex so
+# quoted version specifiers stay in one piece. Unset means do nothing.
+POST_INSTALL_PIP_ARGS = os.environ.get("POST_INSTALL_PIP_ARGS")
+
 
 @session(
     python=get_supported_python_versions(),
@@ -89,6 +96,11 @@ def smoke_tests(session: Session) -> None:
     # required on Python 3.12+ where distutils was removed from stdlib.
     session.install("setuptools")
 
+    # Last, so it replaces any version the installed package pinned.
+    if POST_INSTALL_PIP_ARGS:
+        session.log(f"Applying POST_INSTALL_PIP_ARGS: {POST_INSTALL_PIP_ARGS}")
+        session.install(*shlex.split(POST_INSTALL_PIP_ARGS))
+
     session.log("Running smoke tests")
 
     # Use run_tests.sh to properly handle --junit and other custom flags

docs/src/_static/custom.css

@@ -77,6 +77,18 @@ html.dark code.docutils {
     font-weight: 600;
 }
 
+/* Allow long dotted FQNs to wrap in the sidebar instead of clipping */
+.sy-sidebar-toc,
+.sy-sidebar-toc ul,
+.sy-sidebar-toc li {
+    overflow: visible;
+}
+
+.sy-lside a {
+    white-space: normal;
+    word-break: break-word;
+}
+
 /* Reset code font in sidebar nav — prevents identifier-titled pages
  * from rendering in monospace */
 .sy-lside a code {

scripts/make/build.py

@@ -6,14 +6,27 @@
 """Build the coreai-opt package.
 
 Usage:
-    build.py          Build a standard wheel from the current version
-    build.py --dev    Build a dev wheel with a PEP 440 .dev version
+    build.py --no-sources    Build a release: the `.dev` suffix stripped, via
+                              `uv build --no-sources` (ignores [tool.uv.sources];
+                              the recommended way to build the publishable
+                              artifact). Called by `make build`.
+    build.py --dev            Build a dev wheel with a timestamped PEP 440 .dev
+                              version. Called by `make build-dev`.
+
+``_about.py`` stores ``latest_released_version`` (the last tagged release) by
+hand; ``__version__`` is computed from it — add one to its last number, add
+``.dev0``. This script computes the version to build from
+``latest_released_version``, not from the on-tree ``__version__`` (which must
+never be treated as already released), writes it into ``_about.py``, builds,
+then restores the file. A repo that uses this one as a submodule (building
+one combined wheel) can add its own extra number to the version with
+``COREAI_OPT_VERSION_EXTENSION``; see
+``scripts/release/release_utils.next_release_base``.
 """
 
 from __future__ import annotations
 
 import argparse
-import os
 import subprocess
 import sys
 from pathlib import Path
@@ -24,43 +37,71 @@
 # internal-only scripts use the explicit `external.scripts.*` form instead.
 from scripts._utils import find_repo_root as _find_repo_root
 from scripts.release.release_utils import (
-    get_dev_release_version,
-    get_package_version,
+    get_dev_version_override,
+    get_version_extension,
+    next_release_base,
+    read_about,
+    resolve_build_version,
+    restore_about,
     write_version,
 )
 
 
-def run_build() -> None:
-    """Run ``uv build`` to produce the wheel and sdist."""
-    print(f"Building package with python (version: {sys.version})...")
-    subprocess.run(["uv", "build"], check=True)
-    print("Build complete! Check dist/ directory")
+def run_build(*, no_sources: bool) -> None:
+    """Run ``uv build`` to produce the wheel and sdist.
+
+    Args:
+        no_sources: Pass ``--no-sources`` to ``uv build``, ignoring
+            ``[tool.uv.sources]`` so the artifact doesn't depend on uv-specific
+            index overrides — used for the publishable release build.
+    """
+    sys.stdout.write(f"Building package with python (version: {sys.version})...\n")
+    command = ["uv", "build"]
+    if no_sources:
+        command.append("--no-sources")
+    # `uv build` writes to the same stdout, but Python block-buffers when stdout
+    # isn't a terminal (CI logs, `make build > file`). Flush first so our lines
+    # don't land after the child's. This also flushes anything buffered earlier,
+    # so it's the only flush the script needs.
+    sys.stdout.flush()
+    subprocess.run(command, check=True)
+    sys.stdout.write("Build complete! Check dist/ directory\n")
 
 
 def _build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(description="Build the coreai-opt package.")
-    parser.add_argument(
+    mode = parser.add_mutually_exclusive_group(required=True)
+    mode.add_argument(
         "--dev",
         action="store_true",
-        help="Build a dev wheel with a PEP 440 .dev version",
+        help="Build a dev wheel with a timestamped PEP 440 .dev version",
+    )
+    mode.add_argument(
+        "--no-sources",
+        action="store_true",
+        help="Build a release via `uv build --no-sources`",
     )
     return parser
 
 
 def main() -> None:
     args = _build_parser().parse_args()
-    if not args.dev:
-        run_build()
-        return
     repo_root = _find_repo_root(Path(__file__))
-    original_version = get_package_version(repo_root)
-    dev_version = os.environ.get("DEV_VERSION") or get_dev_release_version(original_version)
+
+    about = read_about(repo_root)
+    release_base = next_release_base(about.latest_released_version, get_version_extension())
+    build_version = resolve_build_version(
+        release_base,
+        dev=args.dev,
+        dev_version_override=get_dev_version_override(),
+    )
     try:
-        write_version(repo_root, dev_version)
-        print(f"Dev version: {dev_version}")
-        run_build()
+        write_version(about, build_version)
+        sys.stdout.write(f"Version: {build_version}\n")
+        run_build(no_sources=args.no_sources)
     finally:
-        write_version(repo_root, original_version)
+        # Restore the exact bytes read before the build rewrote the version.
+        restore_about(about)
 
 
 if __name__ == "__main__":