Commit Graph

11 Commits

Author SHA1 Message Date
Peter Hawkins 82a347438f Fix pytype failures related to teaching pytype about NumPy scalar types.
PiperOrigin-RevId: 519205179
2023-06-02 18:03:40 +01:00
Peter Hawkins c051e6a51d Fix or ignore some pytype errors related to jnp.ndarray == jax.Array.
PiperOrigin-RevId: 511294746
2023-06-02 18:02:16 +01:00
Jake VanderPlas 6fcb84268e Use jax.tree_util.tree_map in place of deprecated tree_multimap.
The latter is a simple alias of the former, so this change is a no-op.

PiperOrigin-RevId: 461229165
2022-07-24 17:53:28 +01:00
Yazhe Li 86307bf254 Packaging BYOL directory so that it can be distributed with pip install.
PiperOrigin-RevId: 408361356
2021-11-09 22:51:14 +00:00
Peter Hawkins a7d75013c9 Remove users of jax.api.* symbols, in preparation for removing the deprecated jax.api name.
In most cases, use the public jax.* name instead.

PiperOrigin-RevId: 395812511
2021-09-17 17:49:29 +01:00
Rebecca Chen 6fdfc75483 Add missing typing.Optional type annotations to function parameters.
PiperOrigin-RevId: 376229010
2021-06-15 23:52:48 +01:00
Florent Altché 2e866f1937 Switch dataset to the new tf.data.Dataset API.
PiperOrigin-RevId: 369679305
2021-04-28 14:54:15 +00:00
Jake VanderPlas 0e5237df2a Use jax.api.device_put_sharded() in place of private JAX APIs.
PiperOrigin-RevId: 332514384
2020-09-23 16:43:36 +01:00
Florent Altché 8457046b2c Add checkpoints from the ablation study.
PiperOrigin-RevId: 328023346
2020-08-26 16:54:56 +01:00
Florent Altché 63fa5e72d5 Fix initial convolution channels not multiplied by width_multiplier.
PiperOrigin-RevId: 327005191
2020-08-21 08:21:04 +00:00
Clara Huiyi Hu 923ad3cff0 This notebook illustrates the CIFAR-10 experiments in the paper:
PiperOrigin-RevId: 326141025
2020-08-21 08:16:49 +00:00