jax

86fcfbfa - Fix memory leak when no axis is provided to pmap. (#3394)

Fix memory leak when no axis is provided to pmap. (#3394) * Fix memory leak when no axis is provided to pmap. * Work around flake8 false positive. Co-authored-by: Matthew Johnson <mattjj@google.com>