jax

9c1003ae - Add clarification to softmax and log_softmax regarding masked-out elements.