jakevdp
Follow
Pinned
1,733 contributions in the last year
Contribution activity
July 2023
Created 27 commits in 1 repository
Created a pull request in google/jax that received 6 comments
Opened 26 other pull requests in 1 repository
google/jax
22
merged
3
open
1
closed
- Add implementation of jnp.put
- lax.dynamic_slice: avoid negative index correction for unsigned indices
- fix deprecations in core.py
- Allow implicit float8 and int4 promotion under numpy_dtype_promotion='full'
- replace use of has_opaque_dtype
- Rename opaque dtype to extended dtype.
- CI: update jupytext to v0.14.7
- Use standard framework for jax.tree* deprecation.
- Refactor opaque dtype implementation.
- custom prng: introduce mechanism to identify key arrays by dtype
- Re-parameterize jax.random.gamma for better behavior at endpoints
- support np.array(x) where x is a custom pytree with __jax_array__
- mypy: suppress annotation-unchecked notes
- remove stray print statement
- [sparse] support preferred_element_type in dot_general
- Lower jax.numpy matmul functions to mixed-precision dot_general
- jnp.einsum: lower to mixed-precision dot_general when possible
- Lower jax.numpy.dot to mixed-precision dot_general
- type promotion: improve implicit type promotion error for float8 & int4
- Make jit work with custom float inputs
- Remove obsolete jaxlib version checks
- pre-commit: update mypy to most recent version
- Remove several deprecated APIs
- Fix scatter batching rule for scatter_apply
- NumpySignaturesTest: account for 'mean' param to std/var
- Some pull requests not shown.
Reviewed 27 pull requests in 1 repository
google/jax
25 pull requests
- Add kl_div and rel_entr functions
- Allow implicit float8 and int4 promotion under numpy_dtype_promotion='full'
- Update incorrect Dirichlet PDF in documentation
- Add implementation of jnp.put
- make _dot_general_batch_rule handle python builtin numeric types
- Apply pyupgrade --py39-plus.
- custom prng: introduce mechanism to identify key arrays by dtype
- Cloud TPU CI: make sure we update test deps and upgrade protobuf version
- Remove unneeded Python 3.9+ version checks
- Remove obsolete numpy version checks
- [shape_poly] Add support for max0 for symbolic dimensions.
- add details on randomness and numerics to API compatibility docs
-
Add type annotation to
jnp.tensordot(..., axes=...) -
Fix docstring for
jax.debug.{print,callback} - jnp.einsum: lower to mixed-precision dot_general when possible
- make scatter-apply jvp notimplemented (for now...)
- Make jit work with custom float inputs
- pre-commit: update mypy to most recent version
-
document and test
CustomVJPPrimaltype as API symbol - einsum: inf inputs could cause superfluous nan outputs
-
random.split: generalize to optional shape parameter - Note that the MSVC studio 2019 redistributable is required for JAX on…
- Fixed _mapped_axis_size raising an uncaught TypeError
- clean up and parameterize random tests more
- test non-threefry RNGs across both typed and raw uint32 key formats
- Some pull request reviews not shown.
Created an issue in openxla/xla that received 2 comments
Dynamic slice with uint8 index returns incorrect result
Reported in google/jax#16836 Here's a minimal repro: import jax.numpy as jnp from jax import lax arr = jnp.arange(200) i = jnp.uint8(128) print(lax.d…
2
comments
Opened 2 other issues in 1 repository
Answered 9 discussions in 1 repository
google/jax
google/jax
-
Does JAX automatically adjust the tangent output in custom_jvp?
This contribution was made on Jul 27
-
How to implement batching rule for custom cuda op?
This contribution was made on Jul 27
-
Custom JVP with non-differentiable PyTree
This contribution was made on Jul 25
-
How do I rotate a 2D array?
This contribution was made on Jul 24
-
Information regarding transpose solver in custom_linear_solve
This contribution was made on Jul 19
-
Sparse matrix-vector products for solvers
This contribution was made on Jul 14
-
how to efficiently calculate sum of cumulative product on sliding window?
This contribution was made on Jul 7
-
vmap function with NamedTuple as input
This contribution was made on Jul 7
-
Is the full set of possible types which could be arguments to JAX primitives available somewhere?
This contribution was made on Jul 6