mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reverts 1a9c94e6265b40c8f1de1e6d920208d648d70fdd

Browse Source 
PiperOrigin-RevId: 568838127
This commit is contained in:
George Necula 2023-09-27 06:57:47 -07:00 committed by jax authors
parent 1a9c94e626
commit ae81ac9cc2
2 changed files with 5 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/_src/config.py
View File

@ -691,8 +691,8 @@ jax_serialization_version = config.define_int_state(
# Note: bump the default serialization version at least one month after
# we update XlaCallModule to support the new version, so that serialized
# modules are forward compatible with deployed versions of XlaCallModule.
# Version 8 of XlaCallModule is supported since July 21th, 2023.
default=int_env('JAX_SERIALIZATION_VERSION', 8),
# Version 7 of XlaCallModule is supported since July 12th, 2023.
default=int_env('JAX_SERIALIZATION_VERSION', 7),
help=(
'The version number to use for native serialization. This must be '
'within the range of versions supported by the tf.XlaCallModule '

8
jax/experimental/jax2tf/README.md
View File

@ -820,14 +820,12 @@ We list here a history of the serialization version numbers:
* Version 7 adds support for `stablehlo.shape_assertion` operations and
for `shape_assertions` specified in `disabled_checks`.
See [Errors in presence of shape polymorphism](https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#errors-in-presence-of-shape-polymorphism). Supported by XlaCallModule
since July 12th, 2023 (cl/547482522),
available in JAX serialization since July 20th, 2023 (JAX 0.4.14),
and the default since August 12th, 2023 (JAX 0.4.15).
since July 12th, 2023 (cl/547482522) and
available in JAX serialization since July 20th, 2023 (JAX 0.4.14).
* Version 8 adds support for the `jax.uses_shape_polymorphism` module
attribute and enables the shape refinement pass only when the
attribute is present. Supported by XlaCallModule since July 21st, 2023
(cl/549973693), available in JAX since July 26th, 2023 (JAX 0.4.14),
and the default since September 27th, 2023 (JAX 0.4.17).
(cl/549973693) and available in JAX since July 26th, 2023 (JAX 0.4.14).
## Known issues