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 upHighlights
- Pro
1,280 contributions in the last year
Contribution activity
June 2020
Created a pull request in google/jax that received 3 comments
- fix error when doing forward-mode of odeint
- removed stale faq entries
- update version and changelog for pypi
- allow closures for odeint dynamics functions
- add remat to top-level docs
- sketch out all-gather primitive
- avoid value-based error check in random.choice
- add back a full_lower, dropped in #3491
- fix an issue with newer versions of pytype
- remove the lower_fun default multiple_results=True
- pin numpy version in setup.py to avoid warnings
- ignore warnings in jax2tf tests
- fix jet typo: jnp not np
- refactor call primitives, simpler param processing
- use original numpy for `arange` shape calculations
- update tpu readme not to pin jax / jaxlib versions
- fix a bug from #3459
- simplify remat partial eval parameterization
- add systematic tests for vmap-of-pmap
- roll back of #3439 while we debug internal fails
- fix vmap-of-pmap mapped_invars logic bug
- add jax.numpy.concatenate(..., axis=None) support
- fix and better explain complex JVPs / VJPs
- print warning when doing jit-of-pmap
- add docstring to ravel_pytree
- Some pull requests not shown.
- jaxpr typechecking: custom rules and errors
- Fix lazy broadcast issue
- de-duplicate constants when lowering to XLA
- Remove unnecessary static_argnum in np.gradient
- Add jax.scipy.signal.detrend
- Cast int8 to bool for lax.not in jax2tf
- Fix typos: np.bool -> np.bool_
- Add class wrapper for double-double arithmetic
- doubledouble: Support repeated doublings
- implement jax.random.choice
- Add experimental precision doubling transform
- add jet primitives, refactor tests
- Jet rule for `integer_pow`
- Propagate raw __name__ of functions wrapped by jit and sharded_jit.
- jaxpr pretty-print: wrap equation RHS when the LHS is long
- add erf_inv rule
- Fix init step sizes to account for the time span of integration
- Document valid enum values for precision.
- Attach source info to Jaxpr equations.
- Initial version of invertible AD implementation
- allow scalar input in poisson sampler
- factor autodiff and vmap tests out from lax_test
- reuse input/output values when joining partially evaluated conditional branches
- Switch CI builds from Travis to Github actions
- Remove special case for scan masking
- Some pull request reviews not shown.
Created an issue in google/jax that received 1 comment
make pmap docs point more clearly to collectives
At the top level of jax.readthedocs.io, we have a section on pmap titled "Parallelism", but that doesn't include any obvious link to the documentat…
1
comment

