[RFC] kwyk pretrained TF->PyTorch weight converter with numerical parity

[dhritimandas]

Summary

What: This adds a converter. It imports the published kwyk TensorFlow
weights into the existing PyTorch kwyk_meshnet architecture. It adds a
harness. The harness checks the converted model against the original
neuronets/kwyk container, end to end.

Why: The kwyk_meshnet architecture (FFG/VWN convolutions, matching
McClure et al. 2019) was already ported to PyTorch in an earlier change.
No PyTorch weights existed for it. This PR fills that one gap only: it
imports the three published TensorFlow checkpoints (MAP, BD, SSD) as
PyTorch state dicts, and validates that the import is numerically correct
against a live copy of the original container -- not just that the
weight-mapping code runs, but that the converted model reproduces the
original's actual output on real data.

This is opened as a draft RFC, not a mergeable-as-is change: two design
questions below need maintainer input before this is final (see "Open
questions").

The converter

nobrainer/datasets/convert_kwyk.py:

• Reads TF variables from a pre-extracted .npz (preferred, no
  TensorFlow needed) or directly from a checkpoint (--tf-path, requires
  tensorflow).
• Maps TF's weight-normalization parameterization (v, g) and axis
  order ([k,k,k,in,out] -> [out,in,k,k,k]) onto FFGConv3d's
  parameters, auto-detecting the TF layer index base (0 vs 1) rather than
  assuming it.
• offline_parity_check independently recomputes each layer's kernel
  from the raw TF v/g and confirms the loaded model reproduces it --
  this gates the weight mapping, before any inference is run.
• Ships as python -m nobrainer.datasets.convert_kwyk --npz ... --n-classes 50 --filters 96 --receptive-field 37 --out kwyk_map.pth, writing a
  .provenance.json sidecar (source checkpoint path + SHA256,
  architecture args) alongside the .pth, since the MAP and BD checkpoints
  are structurally indistinguishable (identical key sets and shapes) and
  the source path is the only durable record of which model a .pth came
  from.

Numerical parity (the part offline_parity_check cannot verify)

offline_parity_check gates the weight mapping only -- it says nothing
about whether the converted model reproduces the original through the
full pipeline (conform -> normalize -> block -> infer -> reassemble) on
real data. scripts/kwyk_reproduction/07_parity_kwyk.py runs both the
converted PyTorch model and a live neuronets/kwyk container on identical
preprocessed input and compares them directly to each other (not each to
ground truth independently, which is a different, weaker check).

Stage A (MAP model, all_50_wn -- the only one of the three
SavedModels that is run-to-run deterministic) compares raw logits on a
real T1 volume (sub-01, via nobrainer.utils.get_data()):

metric	threshold	observed
max|logit diff|	≤ 1e-3	6.87e-05
mean Dice (PyTorch labels vs TF labels)	≥ 0.98	1.0000 (all 49 classes)
voxel agreement	≥ 0.995	1.0000

Ran on a 64^3 crop (8 of 512 blocks), not the full 256^3 volume: the
container is linux/amd64-only and runs under QEMU emulation on the arm64
dev machine this was built on, making a full 512-block run impractically
slow for one sitting. The harness supports full-volume runs unmodified.

Stage B (SSD model, all_50_bvwn_multi_prior) compares MC-aggregate
uncertainty statistics on a 128^3 crop, N=20 samples (full designed scope,
~83 min wall clock under emulation) -- never per-sample, since the two
frameworks' RNGs are unrelated:

metric	threshold	observed
mean-label Dice	≥ 0.95	0.9804
entropy Pearson r	≥ 0.90	0.9963
variance Pearson r	≥ 0.90	0.8740 (misses the gate)

variance_scale_ratio = 0.816: PyTorch's reconciled variance runs ~18%
below TF's in this crop. Plausible, unconfirmed cause: MC sample variance
is a much noisier statistic than the mean at N=20, and PyTorch's
ConcreteDropout3d draws one dropout mask shared across the whole batch
rather than independent per-sample masks -- a real implementation
difference that could suppress inter-sample variance relative to TF's
per-sample sampling. Not silently loosened to force a pass; see the second
open question below.

Open questions for maintainers

1. Where should the converter live? It currently sits in
   nobrainer/datasets/convert_kwyk.py, alongside openneuro.py and
   zarr_store.py -- installed package code, importable as a library. The
   parity harness that validates it lives in scripts/kwyk_reproduction/
   instead -- a standalone, non-packaged reproduction-pipeline directory.
   Should the converter move to sit next to the harness it's validated
   against, or does it belong with the other dataset/conversion utilities
   regardless of where its own validation tooling lives?

2. What parity tolerance is actually acceptable? The thresholds above
   were proposed as initial gates, explicitly meant to be pinned from a
   first real observed run rather than picked in the abstract -- this PR
   is that first run. Stage A's gates all clear with wide margin. Stage
   B's variance gate does not (0.874 vs 0.90). Is 0.90 the right bar given
   N=20 sampling noise on a second-moment statistic, should it be loosened
   with N recorded as a caveat, or should the fix be a larger N (which
   costs real wall-clock time under emulation) before this is called
   validated?

Test plan

• uv run pytest nobrainer/tests/unit/ -m "not gpu" -q -- 427 passed;
  one pre-existing unrelated failure
  (test_croissant.py::test_returns_true_on_valid).
• offline_parity_check passes for all 8 conv layers (incl.
  classifier) on both the MAP and SSD converted checkpoints.
• Real container run: preprocessing gate, Stage A (64^3 crop), Stage B
  (full 128^3/N=20 scope) -- see docs/kwyk_parity_report.md.
• Full 256^3 Stage A run on faster/native (non-emulated) hardware.
• Resolve the two open questions above before taking this out of draft.
• CI