madspace: float32 networks

[theoheimel]

Adds single-precision support to madspace, mainly so the MadNIS networks can be evaluated and trained in float32. This will enable faster MadNIS training and sampling, and reduce gridpack sizes by a factor 2. In terms of unweighting efficiency, i would not expect a difference, but this will need some investigation especially for longer trainings.

Everything else (including normalizing flow jacobians) remains fp64. However, these changes can also be used as a basis to investigate moving more of phasespace to fp32 in the future.

Stacked on #222 (fix-madspace-simd).

Changes

• Data types: DataType is now int32 / float32 / float64 / batch_sizes. This is a clean break: graph JSON, the Python enum and npy files use the new names, and npy and DLPack gain <f4/f32 support. Tensor views now check that the element type matches the tensor dtype.
• Instruction set: an entry can have precision: [float32, float64]. The generator expands it into X_float32 and X_float64, each with its own opcode. The builder method X dispatches on the argument dtype, so frontend code doesn't change. A float64 literal passed to a float32 call is converted first.
• Instructions: matmul (sgemm/dgemm, and the same on cuBLAS/rocBLAS) and all activations exist in both precisions. New convert_float32 and convert_float64 instructions are differentiable.
• Kernels: the CPU and GPU kernel policies are templated on the float type, so float32 and float64 share the kernel code.
• SIMD: float32 vectors have twice the lanes of float64 (NEON 4, AVX2/AVX512Y 8, AVX512 16), with int32 indices and masks of the same width. Jobs are aligned to the larger lane count.
• Training: the global gradient buffer and the Adam step follow the parameter dtype, and all trainable parameters must share one dtype. The gradient clipper computes the norm in float64.
• MLP: MLP(..., precision=...) converts at its boundaries. Flow, DiscreteFlow and ChannelWeightNetwork forward this as subnet_precision / precision.
• MG7: new run-card setting madnis.nn_float_precision (float64 by default, or float32).

Fixes found along the way

• The activation backward kernels overwrote the input gradient instead of adding to it.
• The Python shape check for single (non-batch) inputs compared the wrong dimensions.
• The gather and select backward kernels used float views for int gradients. The new dtype check caught this.
• The bundled OpenBLAS was built without single precision. On Linux this made the doc examples segfault.
• The .pyi stub generation needs to know where DataType lives to render its default arguments.

Testing

• Locally (macOS arm64): all 1789 madspace tests pass in scalar and NEON SIMD modes. The new tests cover:
  • float32 activations and matmul, forward and backward, against torch;
  • the conversion instructions;
  • float32 MLP training against torch for both the built-in Adam and torch Adam;
  • float32 flow outputs and gradients;
  • saving and loading float32 globals.
• MG7 e+ e- > w+ w- with MadNIS: the cross section is 2.6902468 pb in float64 and 2.6902465 pb in float32.
• CI: the CUDA/HIP builds compile. The wheel builds and tests (Linux x86 and macOS arm64, in every SIMD mode the runner supports) and the docs examples pass.

[oliviermattelaer]

fyi

#Should fix

• Backward through stack/cat on float32 fails. The backward of stack, unstack, pop, batch_cat, batch_split and cat still allocates float64 gradients. Their forward passes accept float32, so backward on float32 data throws “invalid dtype in add” or a view dtype error. MadNIS doesn’t hit this today because the networks only use matmul and activations, but any float32 graph built with these ops will. The fix is to allocate with the input’s own dtype.
• Old gridpacks stop loading. Renaming float/int and the NN instructions means gridpacks and graph JSON saved before this PR fail to load. The PR calls this a deliberate clean break, but accepting the old names on load would cost very little.
• GPU tensor_add drops most of a broadcast gradient. This one predates the PR, but the PR rewrote that function without fixing it. When the target is a broadcast gradient (a global’s), only the first event is added, where the CPU version sums the whole batch.

#Minor

• test_flow_float32_subnets leaves fixed-name globals in the shared default context.
• The gradient clipper converts the whole float32 gradient to float64 and back on every step, when only the norm needs float64.
• The float32/float64 branches in the CPU and GPU device ops are near-copies.
• The new dtype_of() isn’t used in two places where it applies.
  nn_precision in launch.py has an error branch that can never run.

[theoheimel]

@oliviermattelaer Thanks, I've fixed the issues (except for the gridpack one, since we're not aiming for cross-version compatibility and are still alpha anyways)