Create your own GitHub profile
Sign up for your own profile on GitHub, the best place to host code, manage projects, and build software alongside 50 million developers.
Sign up
Popular repositories
465 contributions in the last year
Activity overview
Contribution activity
June 2020
Created a pull request in google/jax that received 2 comments
- Revert "Remove documentation presubmit action now that the RTD presubmit is enabled. (#3480)"
- Bump jaxlib version to 0.1.49 and update WORKSPACE
- Add instructions for how to use the TensorBoard profiler to the profiling docs.
- Remove documentation presubmit action now that the RTD presubmit is enabled.
- Allow specifying both `devices` and `axis_size` to pmap.
- Add num_partitions argument to sharded_jit
- Update jax version to 0.1.70
- Handle with_sharding_constraint inside arbitrary subjaxprs.
- Fix handling of infeed token inside sharded_jit
- Add note to docs describing how pytree arguments work.
- Fix logsumexp test
- Make CUDA wheels manylinux2010 compliant, add CUDA 11, drop CUDA 9.2
- Update JAX FAQ.rst, jax.device_put
- Improve dtype test coverage for random_test
- Avoid tuple psum in pmean
- disable new jax2tf TPU test
- update tpu readme not to pin jax / jaxlib versions
- add systematic tests for vmap-of-pmap
- print warning when doing jit-of-pmap
- Switch CI builds from Travis to Github actions
- Add note to docs describing how pytree arguments work.
- Cleanup: remove unused imports in tests
- pmap(in_axes=None) of sharded_jit