261 Commits

Author SHA1 Message Date
Peter Hawkins
b49e31a012 Update version numbers after release. 2022-09-28 18:49:22 +00:00
Peter Hawkins
8d8643664c jax/jaxlib 0.3.20 release candidate. 2022-09-28 13:33:52 +00:00
Skye Wanderman-Milne
d028d93983 Update version and changelog for jax 0.3.19 release 2022-09-27 11:00:27 -07:00
Skye Wanderman-Milne
3c0d280bc0 Update version and changelog for jax 0.3.18 release 2022-09-26 12:43:39 -07:00
Jake VanderPlas
0cb233eec9 Add initial jax.Array base class for instance checks & annotation 2022-09-26 07:48:43 -07:00
Peter Hawkins
bcd36d8eb2 Jax and jaxlib 0.3.18 release candidate. 2022-09-26 14:10:57 +00:00
Yash Katariya
da90234cae Delete soft_pmap as it has no users. Please use pjit or xmap if you do want soft_pmap.
`jax.soft_pmap` is undocumented. If it were documented, a deprecation period would have been provided.

PiperOrigin-RevId: 474145090
2022-09-13 15:52:10 -07:00
Peter Hawkins
40c80d7d0a Remove jax._src from JAX namespace.
This is a JAX-internal name and not subject to any deprecation policy. Please avoid the use of JAX-internal functions outside JAX.

PiperOrigin-RevId: 473243243
2022-09-09 07:06:00 -07:00
Roy Frostig
a2ad414e7c mention AOT readiness in changelog 2022-09-02 13:02:25 -07:00
Jake VanderPlas
ee4ea27c3e update version and changelog for pypi 2022-08-31 11:37:09 -07:00
Peter Hawkins
5527966b27 [JAX] Deprecate .to_py() property on arrays. Implement __array__ instead.
.to_py() was something of an accidental export from the JAX array classes. There are other mechanisms to turn a JAX array into a NumPy array, including `np.asarray(x)` and `jax.device_get(x)`. Deprecate this mechanism because it is redundant.

PiperOrigin-RevId: 469984029
2022-08-25 07:28:27 -07:00
jax authors
0a51a5a2ba Merge pull request #11944 from jakevdp:rm-tile
PiperOrigin-RevId: 469539007
2022-08-23 13:16:50 -07:00
Jake VanderPlas
b8fe0ab8b1 Fix JVP rule for lax.pow() 2022-08-23 11:18:43 -07:00
Jake VanderPlas
8378d08fcd Remove deprecated array.tile() method 2022-08-23 07:36:59 -07:00
Sharad Vikram
b0fdf10a63 Apply suggestions from code review
Co-authored-by: Matthew Johnson <mattjj@google.com>
2022-08-18 10:50:50 -07:00
Sharad Vikram
393bca122d Expose pure callback and enable rank polymorphic callbacks 2022-08-17 10:56:42 -07:00
Matthew Johnson
d19e34fa4a delete old remat implementation
moved lowering rule logic from remat_impl.py (now deleted) to ad_checkpoint.py
2022-08-16 23:16:37 -07:00
Sharad Vikram
841a662894 Update version + changelog 2022-08-11 17:11:52 -07:00
Matthew Johnson
be6f6bfe9f set new jax.remat / jax.checkpoint to be on-by-default 2022-08-10 10:29:38 -07:00
Jake VanderPlas
79406757d0 Remove deprecated jax.experimental.optimizers
The new location is jax.example_libraries.optimizers
2022-08-09 08:50:59 -07:00
Peter Hawkins
c735c6bf0e Increase minimum NumPy version to 1.20.
Per NEP 29, support for 1.19 ended on Jun 21, 2022.
2022-08-06 14:51:14 +00:00
Sharad Vikram
e1da5c7a36
Update CHANGELOG.md 2022-08-02 22:28:52 -07:00
Jake VanderPlas
91dbcbf525 Remove deprecated jax.experimental.stax
The new location is jax.example_libraries.stax
2022-08-02 16:50:06 -07:00
Lena Martens
8ca5ecc7f3 Re-land #11498 after internal fixes.
maintain an alias to `jax.tree_util.tree_map` in the top level `jax` module

PiperOrigin-RevId: 463885774
2022-07-28 11:33:34 -07:00
Jake VanderPlas
108376d792 Remove deprecated function jax.tree_util.tree_multimap 2022-07-26 09:37:27 -07:00
George Necula
afa8f5acb4 Remove jax.experimental.loops. See CHANGELOG
PiperOrigin-RevId: 463297399
2022-07-26 03:39:47 -07:00
Jake VanderPlas
bc90743603 Update changelog for jax/jaxlib v0.3.15 release 2022-07-25 09:47:44 -07:00
George Necula
66dc95e2de removes the jax.mask and jax.shapecheck APIs.
PiperOrigin-RevId: 463026577
2022-07-25 01:23:38 -07:00
George Necula
07fcf79324 jax.mask and jax.shapecheck are being deprecated
Issue: #11557
PiperOrigin-RevId: 462315754
2022-07-21 00:09:31 -07:00
Jake VanderPlas
9090dd179d jax.scipy.linalg.solve: deprecate the sym_pos argument following scipy 1.9.0 2022-07-19 13:57:49 -07:00
jax authors
023e6f5955 Copybara import of the project:
--
e1f1e93e0c8b53e62a064b06b56c84a2bfedb911 by Roy Frostig <frostig@google.com>:

maintain an alias to `jax.tree_util.tree_map` in the top level `jax` module

PiperOrigin-RevId: 461146464
2022-07-15 01:23:51 -07:00
Roy Frostig
e1f1e93e0c maintain an alias to jax.tree_util.tree_map in the top level jax module 2022-07-14 11:00:54 -07:00
Jake VanderPlas
ce08a9fc5c Deprecate top-level aliases of jax.tree_util functions 2022-07-07 11:41:46 -07:00
Jake VanderPlas
39b0ff7eb6 jnp.ndarray: raise TypeError for binary operations with builtin collections 2022-06-29 08:22:05 -07:00
Dan F-M
0788d5708a Implementation of jax.scipy.stats.gaussian_kde 2022-06-28 15:17:12 -04:00
Jake VanderPlas
887abbc3b9 jax.test_util: remove deprecated test classes.
JaxTestCase and JaxTestLoader were deprecated in jax v0.3.1, released Feb 2022.
2022-06-27 11:04:50 -07:00
Peter Hawkins
1e29b7b762 Update CHANGELOG.md and setup.py for 0.3.14 release. 2022-06-27 09:38:41 -04:00
jax authors
406a61cf52 Merge pull request #11146 from sshahrokhi:AbortIfNotInitialized
PiperOrigin-RevId: 457115405
2022-06-24 16:24:57 -07:00
Shiva Shahrokhi
df8c6263de Change JAX_PLATFORMS to raise an exception when platform initialization fails 2022-06-24 21:54:53 +00:00
Ian McKenzie
0cc2ada432 Fix broken links for moved design_notes folder 2022-06-24 12:18:11 -07:00
Jake VanderPlas
f6476f7a03 jnp.roots: better support for computation under JIT 2022-06-23 14:48:53 -07:00
Sharad Vikram
9bd1bd67e0 Update versions for jax/jaxlib release 2022-06-21 12:57:28 -07:00
carlosgmartin
57b89ba7cb Added scipy.stats.gennorm. 2022-06-14 13:38:24 -04:00
jax authors
b174b7751b Merge pull request #10771 from sshahrokhi:gfilecache
PiperOrigin-RevId: 454692872
2022-06-13 13:58:15 -07:00
Shiva Shahrokhi
498ee6007d Using etils(gfile) to support gcs buckets and file system for persistent compilation caching 2022-06-10 00:17:13 +00:00
Jake VanderPlas
d2f80ef117 [x64] deprecate unsafe type casting in scatter-update operations 2022-06-09 15:21:49 -07:00
Sharad Vikram
c0b47fdf2c Update changelog for named_scope and adds it to the docs 2022-06-09 11:22:44 -07:00
Skye Wanderman-Milne
f86282579e Add jax.default_device to CHANGELOG 2022-06-08 14:00:54 -07:00
Sharad Vikram
143ed40a78 Add collect_profile script 2022-06-03 17:56:17 -07:00
carlosgmartin
ca83a80f95 Added random.generalized_normal and random.ball. 2022-06-03 15:11:29 -04:00