1,571 contributions in the last year
Opened 8 other pull requests in 1 repository
Reviewed 15 pull requests in 1 repository
google/jax 15 pull requests
- support implicit broadcasting in transpose rules
- Enable avals-with-names in xmap
- Use the axis names attached to a primitive when selecting the top trace
- Small updates to abstract eval rules (AWN related)
- Add JVP of xmap
- Simplify named shape handling in jax.core
- Speed up test generation 7x
- Avoid forming unnecessary constants when lowering NumPy indexing expr…
- Fix profiler test, attempts to fix CI OOM problems.
- Use convert_element_type instead of device_put_raw.
- DOC: fix typos.
- DeviceArray.__iter__ returns DeviceArrays, without host sync
- use c++ jit dispatch path for op-by-op execution
- [masking] Remove references to masking.Poly from lax and lax_numpy
- [jax2tf] Add support for shape polymorphism for jax2tf; some refactoring of shape computations in JAX core
Created an issue in google/jax that received 2 comments
F = TypeVar("F", bound=Callable) introduced in #4265 may have made our readthedocs type formatting show "~F" for function types: