Trax — Deep Learning with Clear Code and Speed
https://github.com/google/trax
814 forks.
8,015 stars.
122 open issues.
Recent commits:
- [scipy] Use NumPy functions rather than their deprecated aliases in SciPy's namespace.SciPy 1.12.0 removes a number of functions from the scipy namespace that were reexports of NumPy APIs: https://docs.scipy.org/doc/scipy/release/1.12.0-notes.html#expired-deprecationsThis change updates users to use the NumPy names instead.PiperOrigin-RevId: 652894414, Copybara-Service
- Replace deprecated `jax.tree_*` functions with `jax.tree.*`The top-level `jax.tree_*` aliases have long been deprecated, and will soon be removed. Alternate APIs are in `jax.tree_util`, with shorter aliases in the `jax.tree` submodule, added in JAX version 0.4.25.PiperOrigin-RevId: 634149223, Copybara-Service
- Remove uses of deprecated jax.lax.tie_inThis has been a no-op since jax v0.2.0, and passes the second argument through unchanged. `tie_in` will be deprecated as of jax v0.4.24; see https://github.com/google/jax/pull/19413PiperOrigin-RevId: 603220025, Copybara-Service
- [JAX] Replace uses of jax.devices("cpu") with jax.local_devices(backend="cpu").An upcoming change to JAX will include non-local (addressable) CPU devices in jax.devices() when JAX is used multicontroller-style, where there are multiple Python processes.This change preserves the current behavior by replacing uses of jax.devices("cpu"), which previously only returned local devices, with jax.local_devices("cpu"), which will return local devices both now and in the future.This change is always be safe (i.e., it should always preserve the previous behavior) but it may sometimes be unnecessary if code is never used in a multicontroller setting.PiperOrigin-RevId: 582544519, Copybara-Service
- trax: avoid referencing deprecated jax.random.threefry_2x32PiperOrigin-RevId: 575922997, Copybara-Service