Update references to JAX's GitHub repo#81

copybara-service[bot] wants to merge 1 commit intomaingoogle-research/sofima:mainfrom

test_702886847google-research/sofima:test_702886847Copy head branch name to clipboard