fix(distrax): Tanh NaN clip + MultivariateNormal vmap batch_shape fix - #328

Open

Sumu004 wants to merge 2 commits into

google-deepmind:mainfrom

Sumu004:fix/tanh-bijector-nan-clip

Commits on Jul 11, 2026

fix(tanh): clip inverse to (-1+eps, 1-eps) to prevent NaN in log_prob
Sumu004
committed
04802b9
fix(mvn_from_bijector): make batch_shape robust to vmap-prepended dimensions
Sumu004
committed
e222311