Avoid creating NaNs in binary_cross_entropy_with_logits

[miguelbiron]

Use the "double where trick" with a short one liner solution for handling the cases where NaNs are created.

Tackles #2240.

[github-actions]

Benchmark report

this PR bin-x-ent-fix-nans at 3f6041f1 vs baseline master at f107e24d

- run time:     1 slower, 0 faster
+ compile time: 1 faster

Significant changes (2)

                                ──────── run time ───────     ────── compile time ──────
  benchmark                     baseline   this PR      Δ     baseline   this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────
+ predictive_forward_sampling   729.7 ms  733.8 ms  +0.6%     196.2 ms  146.8 ms  -25.2%
- nuts_logistic_regression        1.06 s    1.12 s  +5.4%       3.34 s    3.48 s   +4.2%

_{Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below.}

Full results

distributions

                                 ───────── run time ────────     ────── compile time ─────
  benchmark                      baseline  this PR         Δ     baseline   this PR      Δ
──────────────────────────────────────────────────────────────────────────────────────────
  biject_to_constraints            4.2 ms   4.1 ms     -1.4%     375.5 ms  366.9 ms  -2.3%
  categorical_log_prob             2.2 ms   2.1 ms     -2.2%      70.2 ms   73.9 ms  +5.2%
  dirichlet_log_prob               692 µs   687 µs     -0.7%     459.0 ms  444.2 ms  -3.2%
  dirichlet_sample                45.1 ms  44.8 ms     -0.5%     875.5 ms  882.0 ms  +0.7%
  gamma_log_prob                   2.1 ms   2.3 ms     +4.8%       2.41 s    2.36 s  -2.1%
  gamma_sample                    20.5 ms  20.3 ms     -1.1%     807.9 ms  847.8 ms  +4.9%
  lkj_cholesky_sample              5.4 ms   5.5 ms     +1.0%       1.19 s    1.20 s  +1.1%
  mixture_same_family_log_prob     2.1 ms   2.1 ms     +1.0%     112.9 ms  105.5 ms  -6.6%
  multivariate_normal_log_prob     280 µs   284 µs     +1.3%     162.1 ms  175.7 ms  +8.4%
  normal_log_prob                  699 µs   795 µs  (+13.8%)      56.8 ms   60.3 ms  +6.1%
  normal_sample                   21.6 ms  21.7 ms     +0.5%     204.4 ms  218.5 ms  +6.9%
  stick_breaking_transform         6.4 ms   6.3 ms     -1.0%     210.7 ms  203.5 ms  -3.4%
  student_t_log_prob               3.1 ms   3.1 ms     +0.3%      80.7 ms   81.4 ms  +0.9%
  truncated_normal_log_prob        789 µs   818 µs     +3.6%      57.2 ms   54.7 ms  -4.3%

handlers

                                  ───────── run time ─────────     ────── compile time ──────
  benchmark                       baseline   this PR         Δ     baseline   this PR       Δ
─────────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    40.3 ms   41.3 ms     +2.3%       3.96 s    3.91 s   -1.4%
  log_density_hierarchical          3.7 ms    3.6 ms     -0.3%       1.29 s    1.23 s   -4.5%
  nested_handler_stack              1.4 ms    1.4 ms     +1.4%       783 µs    610 µs  -22.1%
  potential_energy_and_grad          22 µs     26 µs  (+19.0%)     107.0 ms  101.4 ms   -5.2%
+ predictive_forward_sampling     729.7 ms  733.8 ms     +0.6%     196.2 ms  146.8 ms  -25.2%
  trace_seeded_model                827 µs    854 µs     +3.3%     570.3 ms  550.9 ms   -3.4%

mcmc

                             ──────── run time ───────     ────── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR       Δ
────────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    720.4 ms  737.4 ms  +2.4%       3.21 s   3.58 s  +11.4%
  nuts_dense_mass_funnel       1.17 s    1.19 s  +2.1%       2.68 s   2.91 s   +8.5%
  nuts_eight_schools           1.17 s    1.19 s  +1.5%       2.57 s   2.76 s   +7.4%
  nuts_hierarchical_glm        4.95 s    4.92 s  -0.6%       4.92 s   5.09 s   +3.5%
- nuts_logistic_regression     1.06 s    1.12 s  +5.4%       3.34 s   3.48 s   +4.2%
  nuts_vectorized_chains       2.49 s    2.53 s  +1.8%       3.13 s   3.03 s   -3.0%

svi

                                             ──────── run time ───────     ────── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 304.4 ms  305.1 ms  +0.2%       3.08 s   3.39 s   +9.9%
  svi_autodiagonalnormal_hierarchical          1.05 s    1.05 s  +0.0%       5.24 s   5.78 s  +10.2%
  svi_automultivariatenormal_eight_schools   742.4 ms  737.2 ms  -0.7%       4.10 s   4.18 s   +2.1%
  svi_autonormal_logistic                    755.5 ms  770.2 ms  +1.9%       3.41 s   3.53 s   +3.4%
  svi_multi_particle_elbo                      1.51 s    1.54 s  +1.8%       3.54 s   3.69 s   +4.5%
  svi_trace_mean_field_elbo                    1.35 s    1.32 s  -1.9%       5.41 s   5.78 s   +6.8%

Methodology and environment

Each benchmark is set up untimed, then called once with the JAX caches cleared and several more times warm. Run is the fastest warm call; compile is the first call minus that, i.e. the tracing, lowering and XLA compilation the warm calls did not have to pay for.

Both refs were measured on the same runner over 2 interleaved round(s), taking the best observation per benchmark. A result is called neutral when it moves less than ±5% (run) or ±25% (compile), or when the measurement itself is under 1 ms (run) / 50 ms (compile) — a shared CI runner cannot resolve changes below that. Compile time gets the looser band because it is measured once per round rather than best-of-N, and swings by roughly 20% even between two runs of identical code. A delta shown in parentheses did clear its threshold, but on a measurement below the resolution floor, so it is reported without being called a change.

	baseline	this PR
ref	master	bin-x-ent-fix-nans
commit	f107e24d	3f6041f1
numpyro	0.21.0	0.21.0
jax	0.11.0	0.11.0
backend	cpu	cpu
python	3.14.7	3.14.7

Runner: Linux-6.17.0-1020-azure-x86_64-with-glibc2.39, 4 CPUs.

_{Produced by this benchmark run.}

[Qazalbash]

Thank @miguelbiron.

[miguelbiron]

My pleasure @Qazalbash -- big fan of NumPyro, so I need to give back whenever I can.

[miguelbiron]

Pretty freaky error, don't know what could cause this. Maybe just retry?

The hosted runner lost communication with the server. Anything in your workflow that terminates the runner process, starves it for CPU/Memory, or blocks its network access can cause this error.

[Qazalbash]

Pretty freaky error, don't know what could cause this. Maybe just retry?

The hosted runner lost communication with the server. Anything in your workflow that terminates the runner process, starves it for CPU/Memory, or blocks its network access can cause this error.

It is related to GitHub Actions! I am re-triggering the workflow.