Skip to content

fix: import torchjd fails without numpy on Python < 3.14 #778

Description

@ValerianRey

Summary

On Python 3.10 to 3.13, import torchjd fails with NameError: name 'np' is not defined when numpy is not installed. This means that the default installation (pip install torchjd, which only depends on torch) cannot be imported at all, even for features that have nothing to do with numpy (e.g. torchjd.scalarization, torchjd.aggregation.Mean, torchjd.autojac.backward).

This affects every release since v0.12.0 (introduced by #683), including the latest v0.17.0, and current main (30db018).

Reproduction

uv venv --python 3.12 venv
VIRTUAL_ENV=venv uv pip install torchjd torch --extra-index-url https://download.pytorch.org/whl/cpu
venv/bin/python -c "import torchjd"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File ".../site-packages/torchjd/__init__.py", line 5, in <module>
    from .autojac import backward as _backward, mtl_backward as _mtl_backward
  File ".../site-packages/torchjd/autojac/__init__.py", line 10, in <module>
    from ._jac_to_grad import jac_to_grad
  File ".../site-packages/torchjd/autojac/_jac_to_grad.py", line 8, in <module>
    from torchjd._linalg import compute_gramian
  File ".../site-packages/torchjd/_linalg/__init__.py", line 1, in <module>
    from ._dual_cone import DualConeProjector, QuadprogProjector, projector_or_default
  File ".../site-packages/torchjd/_linalg/_dual_cone.py", line 57, in <module>
    class QuadprogProjector(_WithOptionalDeps, DualConeProjector):
  File ".../site-packages/torchjd/_linalg/_dual_cone.py", line 116, in QuadprogProjector
    def _project_weight_vector(self, u: np.ndarray, G: np.ndarray) -> np.ndarray:
                                        ^^
NameError: name 'np' is not defined

Results for an editable install of main with only torch installed:

Python import torchjd
3.10 ❌ NameError
3.11 ❌ NameError
3.12 ❌ NameError
3.13 ❌ NameError
3.14 ✅ OK

Root cause

In src/torchjd/_linalg/_dual_cone.py, numpy is imported as an optional dependency:

with contextlib.suppress(ImportError):
    import numpy as np
    from qpsolvers import solve_qp

But QuadprogProjector._project_weight_vector (and the module-level _to_array) use np.ndarray in their signatures. Before Python 3.14, function annotations are evaluated eagerly when the def statement runs, i.e. at class-definition time, i.e. at import time. With numpy missing, np is unbound, so the module fails to import. Since torchjd/__init__.py → autojac → _jac_to_grad → _linalg → _dual_cone, the whole package fails to import.

Python 3.14 defers annotation evaluation (PEP 649/749), which is why it works there.

The other modules using the same optional-import pattern are fine:

  • aggregation/_fairgrad.py and aggregation/_nash_mtl.py already have from __future__ import annotations.
  • aggregation/_cagrad.py only uses np inside function bodies, not in annotations.

The _WithOptionalDeps mixin is supposed to raise a nice ImportError at instantiation time, but it never gets a chance because the module fails before that.

Why CI didn't catch it

The options: 'none' job in checks.yml should cover this, but it is masked twice:

  1. It runs on the default Python version (3.14), where annotations are lazy.
  2. It installs the test dependency group, which includes torchvision, which depends on numpy. So numpy is always installed in CI anyway.

Suggested fix

Add from __future__ import annotations at the top of src/torchjd/_linalg/_dual_cone.py, consistently with _fairgrad.py and _nash_mtl.py (alternatively, quote the annotations: "np.ndarray").

I tested this locally on Python 3.10 and 3.13 with only torch installed: import torchjd succeeds, torchjd.scalarization.GeometricMean and torchjd.aggregation.Mean work, and UPGrad() raises the intended error:

ImportError: QuadprogProjector requires ['numpy', 'qpsolvers', 'quadprog'] to be installed. Install them with: pip install "torchjd[quadprog_projector]"

To avoid regressions, we could also add a CI check that runs on the lowest supported Python version, with no options and no dependency group, e.g.:

uv venv --python 3.10 && uv pip install . && uv run python -c "import torchjd; import torchjd.aggregation; import torchjd.autojac; import torchjd.autogram; import torchjd.scalarization"

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    cc: ciConventional commit type for changes to the CI (Github workflows and actions).cc: fixConventional commit type for bug fixes of the actual library (changes to src).package: linalg

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions