Skip to content

Handle subnormal mixture weights without XLA flush-to-zero - #5414

Merged
FlorianPfaff merged 2 commits into
mainfrom
fix/jax-subnormal-mixture-weights
Aug 24, 2026
Merged

Handle subnormal mixture weights without XLA flush-to-zero#5414
FlorianPfaff merged 2 commits into
mainfrom
fix/jax-subnormal-mixture-weights

Conversation

@FlorianPfaff

Copy link
Copy Markdown
Owner

Bug

AbstractMixture validates explicit weights on the host, but then determines nonzero components and normalizes using backend comparisons/reductions. On JAX/XLA, representable positive float64 subnormals can be flushed to zero during those operations. A valid mixture such as weights proportional to [2e-311, 1e-311] is therefore rejected as having no nonzero components; the existing backend fallback also sees a zero max and cannot recover the ratio.

This is inconsistent with the robust XLA-subnormal handling already used by AbstractDiracDistribution.

Fix

  • determine exact zero/nonzero support from the already validated host representation;
  • retain the existing backend normalization path for ordinary weights and the existing square-root scaling path for large finite weights;
  • if backend arithmetic loses the weight scale, normalize the finite nonnegative host values after max-scaling and move only the ordinary-sized probabilities back to the active backend;
  • tolerate floating-point overflow traps in backend reductions so the stable fallback can run.

Regression

Adds a focused test with two positive float64 subnormal weights in a 2:1 ratio. The regression verifies that neither component is pruned and that the resulting weights are [2/3, 1/3].

The branch is based on current main after #5411 and has no overlap with the open PR queue.

@FlorianPfaff
FlorianPfaff merged commit 0de6dca into main Aug 24, 2026
18 of 24 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant