parx 0.1.0__py3-none-any.whl

parx/__init__.py

@@ -0,0 +1,111 @@
+"""
+parx — POLyhedral Activation Region Xplorer
+
+Exactly enumerates the linear regions of ReLU neural networks.
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+
+import numpy as np
+
+from parx import (
+    methods,  # noqa: F401 — trigger method registration
+    viz,  # noqa: F401 — expose parx.viz
+)
+from parx._check import check_julia
+from parx._julia_init import ensure_julia  # noqa: F401
+from parx.analysis import (
+    always_active_neurons,
+    complexity_profile,
+    dead_neurons,
+    neuron_activity_rates,
+    region_size_summary,
+)
+from parx.io import iter_state_dicts
+from parx.io_partition import load_partition, save_partition
+from parx.methods import get_method, list_methods
+from parx.network import extract_features, load_network
+from parx.partition import Partition
+from parx.precompile import precompile
+from parx.region import Region
+from parx.viz import (
+    animate_epochs,
+    animate_epochs_video,
+    plot_feature_embedding,
+    region_palette,
+)
+
+check_julia()
+
+try:
+    __version__ = _pkg_version("parx")
+except PackageNotFoundError:  # pragma: no cover - source tree, not installed
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "__version__",
+    "always_active_neurons",
+    "animate_epochs",
+    "animate_epochs_video",
+    "complexity_profile",
+    "compute_partition",
+    "dead_neurons",
+    "ensure_julia",
+    "extract_features",
+    "iter_state_dicts",
+    "list_methods",
+    "load_network",
+    "load_partition",
+    "neuron_activity_rates",
+    "Partition",
+    "plot_feature_embedding",
+    "precompile",
+    "Region",
+    "region_palette",
+    "region_size_summary",
+    "save_partition",
+]
+
+
+def compute_partition(
+    source,
+    data: np.ndarray,
+    *,
+    method: str = "sparse_julia",
+    include_output_layer: bool = False,
+    **method_kwargs,
+) -> Partition:
+    """Find the linear regions of a ReLU network.
+
+    Parameters
+    ----------
+    source:
+        A single network's parameters.  Accepted forms: PyTorch ``state_dict``,
+        ``nn.Module``, or path to a ``.pth`` / ``.h5`` file.  For per-epoch
+        analyses, iterate over state dicts at the call site (see
+        :func:`parx.io.iter_state_dicts`).
+    data:
+        Input points, shape ``(N, input_dim)``.  Sparse methods scan the array
+        for activation patterns; exact methods use ``data[0]`` as the DFS
+        starting point.
+    method:
+        Name of a registered region-finding method.  Built-ins:
+        ``"sparse_julia"`` (default), ``"exact_julia"``, ``"sparse_python"``,
+        ``"exact_python"``.  See :func:`parx.list_methods`.
+    include_output_layer:
+        Include the final linear layer in the partition.  Defaults to ``False``
+        because only hidden ReLU layers define the polyhedral partition.
+    **method_kwargs:
+        Forwarded verbatim to the chosen method's function.
+
+    Returns
+    -------
+    Partition
+    """
+    weights, biases = load_network(source, include_output_layer=include_output_layer)
+    fn = get_method(method)
+    result = fn(weights, biases, data, **method_kwargs)
+    return Partition.from_result(result, weights, biases)

parx/_check.py

@@ -0,0 +1,14 @@
+import shutil
+
+
+def check_julia() -> None:
+    """Raise a helpful error if Julia is not found on PATH."""
+    if shutil.which("julia") is None:
+        raise RuntimeError(
+            "\n\nparx requires Julia to be installed and available on PATH.\n"
+            "Install Julia via juliaup (recommended):\n\n"
+            "  macOS/Linux:  curl -fsSL https://install.julialang.org | sh\n"
+            "  Windows:      winget install julia -s msstore\n\n"
+            "After installing, restart your terminal and try again.\n"
+            "See https://github.com/Johanmkr/parx for more details.\n"
+        )

parx/_julia_init.py

@@ -0,0 +1,106 @@
+"""
+Handles Julia runtime initialization.
+Import this module once; subsequent calls to ensure_julia() are no-ops.
+"""
+
+import json
+import os
+import shutil
+from pathlib import Path
+
+_julia_initialized = False
+_jl = None
+
+
+def _resolve_juliaup_shim() -> str | None:
+    """Work around a juliapkg bug when Julia was installed via juliaup.
+
+    juliaup puts a launcher shim (not a real Julia binary) at
+    ``~/.juliaup/bin/julia``; the actual per-version installs live under
+    ``~/.julia/juliaup/julia-<version>+.../``.  juliapkg tries to
+    opportunistically upgrade to the newest Julia release on every resolve;
+    when that install attempt fails, it silently falls back to the shim path
+    itself instead of one of the already-installed real binaries. juliacall
+    then derives the system-image path relative to the shim's directory,
+    which never has a ``lib/julia/sys.so`` next to it, and the first Julia
+    call crashes with "could not load library ... sys.so ... No such file or
+    directory".  See CONTRIBUTING.md's Troubleshooting section.
+
+    Returns the absolute path to the real Julia binary juliaup's default
+    channel points at, or ``None`` if ``julia`` isn't on PATH, isn't a
+    juliaup shim, or the real binary can't be resolved for any reason — in
+    which case juliapkg's normal resolution is left alone.
+    """
+    julia_on_path = shutil.which("julia")
+    if julia_on_path is None:
+        return None
+
+    # juliaup's shim is typically a symlink to a binary literally named
+    # "julialauncher"; a plain (non-juliaup) Julia install is not. Some
+    # installs may ship a separate "julia" launcher alongside "julialauncher",
+    # so also accept a sibling julialauncher in the same directory.
+    real_target = os.path.realpath(julia_on_path)
+    launcher_names = ("julialauncher", "julialauncher.exe")
+    if os.path.basename(real_target) not in launcher_names:
+        shim_dir = os.path.dirname(julia_on_path)
+        ext = ".exe" if os.name == "nt" else ""
+        candidate = os.path.join(shim_dir, "julialauncher" + ext)
+        if not os.path.isfile(candidate):
+            return None
+        real_target = candidate
+
+    # juliaup does not follow JULIA_DEPOT_PATH, but defines its own
+    # override for ~/.julia (matches juliapkg's own lookup).
+    depot = os.environ.get("JULIAUP_DEPOT_PATH") or os.path.join(
+        os.path.expanduser("~"), ".julia"
+    )
+    judir = os.path.join(depot, "juliaup")
+    try:
+        with open(os.path.join(judir, "juliaup.json")) as f:
+            meta = json.load(f)
+        version = meta["InstalledChannels"][meta["Default"]]["Version"]
+        info = meta["InstalledVersions"][version]
+        if "BinaryPath" in info:
+            exe = os.path.join(judir, info["BinaryPath"])
+        else:
+            ext = ".exe" if os.name == "nt" else ""
+            exe = os.path.join(judir, info["Path"], "bin", "julia" + ext)
+        exe = os.path.abspath(exe)
+    except (OSError, KeyError, ValueError, json.JSONDecodeError):
+        return None
+
+    return exe if os.path.isfile(exe) else None
+
+
+def ensure_julia():
+    """Initialize the Julia runtime and load the LinearRegions module.
+
+    juliacall/juliapkg manages its own Julia project environment, so we load
+    our LinearRegions module via include() rather than registering it as a
+    package.  Safe to call multiple times — only runs once per process.
+    """
+    global _julia_initialized, _jl
+
+    if _julia_initialized:
+        return _jl
+
+    os.environ.setdefault("JULIA_NUM_THREADS", "auto")
+    # Let Julia own signal handling so its GC threads don't conflict with
+    # Python's signal machinery.  Must be set before juliacall is imported.
+    os.environ.setdefault("PYTHON_JULIACALL_HANDLE_SIGNALS", "yes")
+
+    # Work around the juliaup-shim juliapkg bug (see _resolve_juliaup_shim's
+    # docstring).  setdefault so an explicit user override always wins.
+    juliaup_exe = _resolve_juliaup_shim()
+    if juliaup_exe is not None:
+        os.environ.setdefault("PYTHON_JULIAPKG_EXE", juliaup_exe)
+
+    from juliacall import Main as jl
+
+    julia_file = Path(__file__).parent / "julia" / "src" / "LinearRegions.jl"
+    jl.seval(f'include("{julia_file.as_posix()}")')
+
+    _jl = jl
+    _julia_initialized = True
+
+    return _jl

parx/_lp.py

@@ -0,0 +1,59 @@
+"""Shared LP primitives used by region finders and partition verification."""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.optimize import linprog
+
+_ZERO_NORM_TOL = 1e-10
+_SLACK_TOL = 1e-8
+
+
+def chebyshev_center(
+    D: np.ndarray,
+    g: np.ndarray,
+    *,
+    max_radius: float = 1e3,
+) -> tuple[np.ndarray | None, float]:
+    """Chebyshev centre of ``{x : D x ≤ g}``.
+
+    Returns ``(x_interior, radius)`` or ``(None, 0.0)`` when empty / degenerate.
+
+    Zero-norm rows are filtered (they would make the LP unbounded); a
+    zero-norm row with strictly negative ``g[i]`` makes the system infeasible
+    (returns ``(None, 0.0)``).  Unbounded polytopes are detected by the radius
+    hitting ``max_radius``; callers can compare against that value to decide
+    whether to treat the region as unbounded.
+    """
+    m, n = D.shape
+    if m == 0:
+        return np.zeros(n), float("inf")
+
+    row_norms = np.linalg.norm(D, axis=1)
+    valid = row_norms > _ZERO_NORM_TOL
+
+    if np.any(g[~valid] < -_ZERO_NORM_TOL):
+        return None, 0.0
+    if not valid.any():
+        return np.zeros(n), float("inf")
+
+    D_v = D[valid]
+    g_v = g[valid]
+    norms_v = row_norms[valid]
+
+    # Variables = (x_1, …, x_n, r).  Constraint i: D_v[i] · x + ‖D_v[i]‖ · r ≤ g_v[i].
+    A_ub = np.hstack([D_v, norms_v[:, None]])
+    b_ub = g_v
+    c = np.zeros(n + 1)
+    c[-1] = -1.0  # maximise r ↔ minimise -r
+    bounds = [(None, None)] * n + [(0.0, max_radius)]
+
+    res = linprog(c, A_ub=A_ub, b_ub=b_ub, bounds=bounds, method="highs")
+    if not res.success or res.x is None:
+        return None, 0.0
+
+    x = res.x[:n]
+    r = float(res.x[-1])
+    if r < _SLACK_TOL:
+        return None, 0.0
+    return x, r