[Compatibility Issue] Missing numpy>=2 constraint — np.typing crashes training entry points on numpy<2

[xyf5432]

Affected File

ffn/training/examples.py:130

Current Code

# ffn/training/examples.py
# (no from __future__ import annotations)

def update_seeds(self, batched_seeds: np.typing.ArrayLike):
    """Distributes updated predictions back to the generator buffers.
    ...
    """
    assert self._seeds is not None
    batched_seeds = np.asarray(batched_seeds)

No from __future__ import annotations — the annotation is evaluated when the BatchExampleIter class body executes at module import time.

Root Cause

numpy.typing is only accessible as a module attribute (np.typing.XXX) in numpy >= 2.0.0 (via __getattr__). On numpy < 2.0, access without a prior import numpy.typing triggers:

AttributeError: module 'numpy' has no attribute 'typing'. Did you mean: '_typing'?

Dependency

setup.py declares 'numpy>=1.11.1' (2020-era floor, no upper bound). No lock files exist. Fresh pip installs resolve to numpy 2.x (fine), but any environment with a pre-existing numpy<2 pin (e.g. from an older tensorflow install) keeps numpy 1.x and crashes.

Upstream Dependency Protection — Could Be Avoided, But Not Guaranteed

If any code in the import chain executes import numpy.typing before examples.py loads, the annotation evaluates fine even on numpy<2 (the import injects typing into numpy.__dict__). So the crash may be avoided depending on what the user imports first. However, this protection is incidental — it depends on the internal implementation of dependencies, which can change — so it should not be relied upon.

What we have checked so far:

• pandas: declared in setup.py, but only imported in ffn/utils/decision_point.py, which is loaded only by a test. On the runtime training path, pandas is never imported before examples.py. (Note: only pandas >= 3.0 triggers numpy.typing at import time; pandas 2.x does not.)
• xarray / zarr: not dependencies.
• scipy: does not trigger numpy.typing at import scipy time (verified).
• tensorflow, connectomics (git dep), absl, ffn.training.augmentation and the rest of the training import chain: not yet checked one by one — any of them could potentially do import numpy.typing and avoid the crash.

Impact

ffn/training/examples.py is loaded eagerly by both training entry points:

• python train.py — from ffn.training import examples at line 33
• JAX path — ffn.jax.main → ffn.jax.train → ffn.jax.input_pipeline

import ffn itself is safe (zero imports in ffn/__init__.py), and the inference path does not load examples.py.

On numpy<2 with no prior import numpy.typing in the process, the module raises:

AttributeError: module 'numpy' has no attribute 'typing'. Did you mean: '_typing'?

(the failure mode itself was verified on numpy 1.26.4; the full python train.py reproduction with all dependencies installed has not been run).

Solution

Any of the following in examples.py:

1. from __future__ import annotations — defers annotation evaluation (simplest, also a free import-time win)
2. import numpy.typing before the class definition
3. Use np.ndarray or a plain type in the annotation — the current ArrayLike annotation is also type-incorrect: the docstring says it accepts "array-like object backed by accelerator memory" (i.e. JAX arrays), which np.typing.ArrayLike does not describe.

Optionally bump the declared floor in setup.py from numpy>=1.11.1 to numpy>=2 (or at least >=1.20).

References

• NumPy 2.0 release notes — numpy.typing exposed as module attribute via __getattr__
• Similar fix: numexpr#540