[uermel]

Adds an opt-in --xla flag to copick inference easymode that runs each model's forward pass as one XLA-compiled function, tf.function(jit_compile=True). It is the option released upstream as easymode segment --xla in easymode 1.2.6 (mgflast/easymode#37) and ais segment --xla in ais-cryoet 1.2.40 (mgflast/Ais#16). This PR also moves the supported easymode release from 1.2.5 to 1.2.6.

--xla

• core/xla.py: XLAModel wraps a Keras model.
  • It is called like the model and hands every other attribute to the model. easymode's tiling of an .h5 network and the .scnm inference use it unchanged, and ScnmModel.is_slab still reads the network's shapes.
  • The first call of each input shape compiles, and the compile time is logged.
  • If a compiled call raises, the worker logs one warning and runs the plain model from then on.
  • This package keeps its own wrapper rather than easymode's XLAModel for three reasons: the attribute pass-through, the log line per compiled shape, and the fallback reason in the report.
• InferenceSettings.xla, dispatch_easymode_inference(xla=...) and run_easymode_inference(xla=...) carry the option to every worker. load_stages wraps the network when it is set.
• CLI: --xla/--no-xla, off by default. It stays opt-in because a missing or too old ptxas can abort TensorFlow outright, which no fallback can catch.
• Report: a top-level "xla", and per worker xla_fallbacks, each with feature and reason. A fallback is not an error, and its runs still count as processed.

easymode 1.2.6

• EASYMODE_VERSION / EASYMODE_COMMIT now point to tag v1.2.6 (00ab903), and the install line follows them. The .h5 path uses only _pad_volume, _segment_tomogram_instance and load_model, and all three are unchanged in 1.2.6.
• The .scnm port now follows ais-cryoet 1.2.40. Its inference differs from 1.2.38 in only three ways:
  • XLA wrapper: it adds one, which this package already has.
  • Slab --stride: a new option whose default stays at half the slab depth. easymode does not pass it, so nothing changes here.
  • Last-batch padding under XLA: a short last 2.5D batch is now padded with zeros to the batch size, so the compiled function keeps one input shape. This is ported, although at the .scnm batch size of 1 it never fires.
• The docs name easymode's current metadata pins: numpy<2, tensorflow<2.16 and ais-cryoet, which needs Keras 2. They also drop the claim that easymode is not on PyPI. The install stays --no-deps.

Measurements

Three 10521 tomograms, each 320 x 772 x 548 voxels, streamed from the CryoET Data Portal. TTA 4 on one L40S, with TensorFlow 2.21 and Keras 3.15:

model	plain, per tomogram	XLA, per tomogram	compile	job wall	mask IoU vs plain
microtubule sv3-3d, .h5	310-332 s	106-107 s, 133 s for the first	26.8 s	1002 s to 393 s	0.99994 or better
actin sv2-3d, .scnm slab	86-92 s	67 s, 75 s for the first	7.8 s	292 s to 243 s	0.99982 or better

• No fallback occurred. The 2.5D .scnm slice models were not measured.
• Repeating --xla with easymode 1.2.6 on one of the tomograms reproduced the masks:
  • actin: identical;
  • microtubule: 9 voxels out of 135 million differ, IoU 0.99995. That rerun was on a different GPU.

Tests: new unit tests cover the wrapper, the last-batch padding, the flag and the report. The suite passes with easymode 1.2.6 inside the picking image: 75 passed.

🤖 Generated with Claude Code