PyAETNA 4.0.0__py3-none-any.whl

PyAETNA/__init__.py

@@ -0,0 +1,9 @@
+"""PyAETNA: JAX-accelerated modelling of exoplanet light curves and radial velocities."""
+
+from ._xla_flags import apply_default_xla_flags
+from .version import __version__
+
+# Before JAX initialises its GPU backend (see PyAETNA._xla_flags).
+apply_default_xla_flags()
+
+__all__ = ["__version__"]

PyAETNA/__main__.py

@@ -0,0 +1,5 @@
+"""``python -m PyAETNA config.yaml``: same as the ``pyaetna`` command."""
+
+from .cli import main
+
+main()

PyAETNA/_device_config.py

@@ -0,0 +1,57 @@
+"""Configure JAX device count and BLAS threads for CPU parallelism.
+
+This module MUST be imported before any JAX operations to enable multi-core
+CPU parallelism. It sets XLA_FLAGS to create virtual devices for parallel
+MCMC chains.
+
+IMPORTANT: This configuration ONLY affects the CPU backend. GPU devices are
+still detected and preferred if available. The --xla_force_host_platform_device_count
+flag does not prevent JAX from using GPUs when present on the system.
+"""
+
+import os
+
+# Configure host CPU device count for JAX before importing it.
+# Use the machine's CPU count, and at least _MIN_HOST_DEVICES, by default, but
+# allow explicit override via env var PYAETNA_XLA_FORCE_HOST_PLATFORM_DEVICE_COUNT.
+# Nested sampling on a CPU gives each of its slice chains (50 by default) its own
+# device, which is fastest even when there are fewer cores than devices
+# (sampling/nested_sampling.py, cpu_chain_mesh).
+# NOTE: This setting does NOT affect GPU detection or usage.
+_MIN_HOST_DEVICES = 64
+_cpu_count = os.cpu_count() or 1
+_default = max(_cpu_count, _MIN_HOST_DEVICES)
+_override = os.environ.get("PYAETNA_XLA_FORCE_HOST_PLATFORM_DEVICE_COUNT")
+try:
+    _n_devices = max(1, int(_override)) if _override is not None else _default
+except ValueError:
+    _n_devices = _default
+
+# One thread per BLAS/LAPACK call, unless the user set it. On a CPU, JAX runs dense linear algebra (Cholesky
+# factorisations and triangular solves of dense GPs) as LAPACK calls into OpenBLAS; with many host devices
+# calling it at once, a multithreaded OpenBLAS oversubscribes the CPU and its threads contend. On the K2-100
+# multi-GP benchmark (8 logical CPUs, 50 chains on host devices) one thread made nested sampling 65 times
+# faster (98 -> 6.4e3 likelihood evaluations per second). Must be set before NumPy is first imported.
+for _var in ("OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "OMP_NUM_THREADS"):
+    os.environ.setdefault(_var, "1")
+
+_flag = f"--xla_force_host_platform_device_count={_n_devices}"
+_existing = os.environ.get("XLA_FLAGS", "")
+if "--xla_force_host_platform_device_count=" not in _existing:
+    os.environ["XLA_FLAGS"] = (_existing + " " + _flag).strip()
+
+# Longer timeouts for collectives between host devices, unless the user set them. The slice chains of a
+# nested-sampling iteration meet at an all-gather; XLA aborts the process (SIGABRT) when a device arrives more
+# than ~60 s after the others (warning after 20 s). With 50 chains on few cores and an expensive likelihood a
+# chain can legitimately be that late, so warn after 5 min and abort only after an hour (an interrupted
+# nested-sampling run resumes from its checkpoint). Note (2026-09-30, under investigation): on the laptop, TOI-216
+# (a GP on 196k points, 50 chains on 12 cores) twice stopped with two devices that never reached the all-gather
+# while the CPU was idle, i.e. a deadlock, not a late chain; the same run on 12 cores of the server completed.
+_COLLECTIVE_TIMEOUTS = {
+    "--xla_cpu_collective_call_warn_stuck_timeout_seconds": "300",
+    "--xla_cpu_collective_call_terminate_timeout_seconds": "3600",
+}
+for _name, _value in _COLLECTIVE_TIMEOUTS.items():
+    _existing = os.environ.get("XLA_FLAGS", "")
+    if f"{_name}=" not in _existing:
+        os.environ["XLA_FLAGS"] = f"{_existing} {_name}={_value}".strip()

PyAETNA/_xla_flags.py

@@ -0,0 +1,50 @@
+"""XLA compiler flags PyAETNA needs on GPU.
+
+Applied when the package is imported, i.e. before JAX initialises its GPU
+backend, which is when ``XLA_FLAGS`` is read.  A flag the user already set in
+``XLA_FLAGS`` is never overridden.
+
+``--xla_gpu_experimental_enable_fusion_autotuner=false``
+    XLA's GPU *fusion* autotuner can explode on some fused kernels once the
+    batch is large enough: vectorising the Pi Men likelihood (four
+    instruments, 223 RVs) over 500 points drove compilation above 23 GB of
+    host RAM until the kernel killed the process.  The threshold depends on
+    the model and the code version: with JAX 0.11.0 and the eccentric Pi Men
+    model, 302 points compile in 1.9 s with 1.7 GB, while 303 were killed at an
+    8 GB cap after 64 s.  Nested sampling evaluates its 500
+    initial live points in one batch, so it hit this directly.  Disabling
+    only the fusion autotuner removes the problem and leaves GEMM autotuning
+    on: on K2-140 and on a dense GP the run time is unchanged (turning off
+    all autotuning with ``--xla_gpu_autotune_level=0`` would make a
+    small-batch dense GP ~30% slower).  To re-enable it, set the flag to
+    ``true`` in ``XLA_FLAGS`` before importing PyAETNA.  Reported upstream as
+    https://github.com/openxla/xla/issues/49549 (still present in JAX 0.11.2),
+    with a standalone reproducer; possibly related: tue-bmd/zea#599 (compile
+    time exploding above a grid size, same workaround).  Fixed upstream on
+    2026-09-29 by XLA commit c644af1 ("Prevent exponential op duplication in
+    TritonXLAFoldTransposePass"): transpose folding cloned the operations of
+    reconvergent elementwise graphs into a tree, so compile time and memory
+    grew exponentially.  Keep the flag until the minimum supported JAX
+    includes that commit, then check the reproducer and remove it.
+"""
+
+from __future__ import annotations
+
+import os
+
+_DEFAULT_FLAGS = {
+    "--xla_gpu_experimental_enable_fusion_autotuner": "false",
+}
+
+
+def apply_default_xla_flags(environ: dict[str, str] | None = None) -> None:
+    """Append PyAETNA's default XLA flags unless the user already set them."""
+    env = os.environ if environ is None else environ
+    existing = env.get("XLA_FLAGS", "")
+    extra = [
+        f"{flag}={value}"
+        for flag, value in _DEFAULT_FLAGS.items()
+        if f"{flag}=" not in existing
+    ]
+    if extra:
+        env["XLA_FLAGS"] = " ".join([existing, *extra]).strip()