Tags: vfdev-5/flax
Tags
Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` u… …seful. Fixes google#4008 PiperOrigin-RevId: 646679331
Merge pull request google#3936 from google:nnx-stabilize PiperOrigin-RevId: 636951986
Merge pull request google#3876 from google:nnx-v0.1 PiperOrigin-RevId: 628712571
Fix tests after applying JAX key-reuse checker. See: - https://jax.readthedocs.io/en/latest/jax.experimental.key_reuse.html#experimental-key-reuse-checking - https://jax.readthedocs.io/en/latest/_autosummary/jax.random.clone.html#jax.random.clone PiperOrigin-RevId: 615600399
Merge pull request google#3682 from IvyZX:push PiperOrigin-RevId: 605087622
Merge pull request google#3644 from chiamp:release PiperOrigin-RevId: 600921456
-
Jan 23, 2024 - 2a84d42 79F7
- zip
- tar.gz
Merge pull request google#3444 from levskaya:rc7.5 PiperOrigin-RevId: 577359040
Merge pull request google#3331 from levskaya:version_fix PiperOrigin-RevId: 565109475
Merge pull request google#3326 from IvyZX:push PiperOrigin-RevId: 564889489
PreviousNext