Preserve slow EMA decay in low-precision updates

[sylvesterkaczmarek]

Summary

Compute decay complements, floating-point EMA updates and debiasing in at least float32, then restore the existing stored-state and output dtypes. This prevents valid slow decays from rounding to one in float16/bfloat16. State keys, scan carry types, explicit wider initialization and existing integer-input behavior are retained.

Low-precision checkpoint state remains low precision, so ordinary accumulation rounding is not eliminated. This change is independent of the documentation-only edits in PR #888.

Reproduction

A decay of 0.9999 rounds to one when converted to float16 or bfloat16. ExponentialMovingAverage then returns NaNs for the first debiased update, or a zero non-debiased update. A constant input [1, 2] should instead have first debiased result [1, 2].

Validation

python -m pytest -q haiku/_src/moving_averages_test.py haiku/_src/moving_averages_precision_test.py

35 tests passed. Twenty-one new cases cover three dtypes and decay values, debiased/non-debiased scans against a quantized-state reference, unchanged preview state, warmup, parameter trees, explicit float32 initialization, mixed input/state dtypes and integer compatibility. The complete existing moving-average test module passes. Existing deprecated initialization warnings remain. Persistent-state precision and schema are not changed.

Negative control on unchanged main: 11 failed, 10 passed in 1.72s.

Tested on macOS CPU with real module imports. New-test formatting, scoped static checks, syntax checks and git diff --check pass. GPU/TPU and the full repository suite were not run. No dependency or workflow changes.