jax

4b1bb189 - Fix logsumexp test (#3563)

Fix logsumexp test (#3563) Previously, this test could generate an axis out of bounds error with large num_generated_cases. (discovered in the process of testing #3561)