jax

6ca46723 - Merge pull request #20578 from carlosgmartin:logsumexp_where