mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
13,640 Commits 503 Branches 263 Tags
Commit Graph

281 Commits

Author SHA1 Message Date
Yash Katariya
cc5af7ed98 Rename ReshapeableDevicesSharding to PositionalSharding and add an alias NamedSharding for MeshPspecSharding. 
`MeshPspecSharding` name will be replaced with `NamedSharding` in 3 months.

PiperOrigin-RevId: 485753078
 2022-11-02 19:13:13 -07:00
jax authors
ef63f75e39 Merge pull request #13039 from skye:cache_compile_time_heuristic 
PiperOrigin-RevId: 485644419
 2022-11-02 11:13:52 -07:00
Skye Wanderman-Milne
cc5171034f Add new config jax_persistent_cache_min_compile_time_secs. 
This replaces `jax_persistent_cache_min_instruction_count` introduced
in https://github.com/google/jax/pull/12798, since gating on the
compile time seems strictly better than gating on the instruction
count (except maybe that the instruction count is more deterministic,
but I don't think that's a big deal).

I defaulted to 1 second as the minimum threshold based on the same
flax wmt example
(https://github.com/google/flax/tree/main/examples/wmt) numbers from

name | instruction_count | compile_time_secs
---- | ----------------- | -----------------
`broadcast_in_dim` | 2 | 0.01633763313
`convert_element_type` | 2 | 0.01704716682
`reshape` | 2 | 0.01730203629
`_squareit` | 2 | 0.01730823517
`broadcast_in_dim` | 2 | 0.0182030201
`convert_element_type` | 2 | 0.01982188225
`concatenate` | 2 | 0.02102327347
`true_divide` | 2 | 0.02172231674
`broadcast_in_dim` | 2 | 0.02370619774
`broadcast_in_dim` | 2 | 0.02393102646
`broadcast_in_dim` | 2 | 0.02488565445
`broadcast_in_dim` | 2 | 0.03395628929
`broadcast_in_dim` | 2 | 0.03428125381
`broadcast_in_dim` | 2 | 0.0394551754
`shift_right_logical` | 2 | 0.06500506401
`<lambda>` | 3 | 0.01793265343
`_unstack` | 5 | 0.01975226402
`_reduce_sum` | 5 | 0.0210878849
`_reduce_sum` | 5 | 0.02416801453
`_multi_slice` | 9 | 0.09065580368
`_threefry_split` | 232 | 0.09037566185
`_threefry_split` | 232 | 0.09161829948
`<unnamed wrapped function>` | 2668 | 7.701903343
`<unnamed wrapped function>` | 3455 | 17.57672167
`<unnamed wrapped function>` | 46580 | 166.2570884
`init` | 60361 | 26.35722399
`<unnamed wrapped function>` | 78010 | 3.879326344

Also adds new float config functionality.
 2022-11-02 00:56:19 +00:00
Jake VanderPlas
2416d15435 Call _check_arraylike for jnp.linalg & jnp.fft functions 2022-10-31 09:19:53 -07:00
Peter Hawkins
bf21391248 [JAX] Change the default pmap() ordering to match the ordering of jax.devices() for single-process TPU jobs. 
PiperOrigin-RevId: 484062717
 2022-10-26 13:56:07 -07:00
Peter Hawkins
ce9e009c4c [JAX:CPU] Enable buffer donation on CPU. 
Fix a bug in PJRT where if a buffer was not owned (e.g., it aliased a NumPy buffer) it could still be donated and that would lead to a use after free.

PiperOrigin-RevId: 484001545
 2022-10-26 10:13:01 -07:00
Jake VanderPlas
2009e65a33 jnp.gradient: call check_arraylike on inputs & clean-up implementation 2022-10-24 15:27:33 -07:00
Jake VanderPlas
4aceb81570 Add docs & changelog for jax.scipy.stats.mode 2022-10-20 15:55:57 -07:00
Skye Wanderman-Milne
81eb3fca55 Add new config jax_persistent_cache_min_instruction_count. 
This can be used to limit the number of entries written to the
persistent compilation cache.

I defaulted to setting 6 as the minimum threshold based on running the
flax wmt example
(https://github.com/google/flax/tree/main/examples/wmt) and logging
the instruction counts and complilation time:

name | instruction_count | compile_time_secs
---- | ----------------- | -----------------
`broadcast_in_dim` | 2 | 0.01633763313
`convert_element_type` | 2 | 0.01704716682
`reshape` | 2 | 0.01730203629
`_squareit` | 2 | 0.01730823517
`broadcast_in_dim` | 2 | 0.0182030201
`convert_element_type` | 2 | 0.01982188225
`concatenate` | 2 | 0.02102327347
`true_divide` | 2 | 0.02172231674
`broadcast_in_dim` | 2 | 0.02370619774
`broadcast_in_dim` | 2 | 0.02393102646
`broadcast_in_dim` | 2 | 0.02488565445
`broadcast_in_dim` | 2 | 0.03395628929
`broadcast_in_dim` | 2 | 0.03428125381
`broadcast_in_dim` | 2 | 0.0394551754
`shift_right_logical` | 2 | 0.06500506401
`<lambda>` | 3 | 0.01793265343
`_unstack` | 5 | 0.01975226402
`_reduce_sum` | 5 | 0.0210878849
`_reduce_sum` | 5 | 0.02416801453
`_multi_slice` | 9 | 0.09065580368
`_threefry_split` | 232 | 0.09037566185
`_threefry_split` | 232 | 0.09161829948
`<unnamed wrapped function>` | 2668 | 7.701903343
`<unnamed wrapped function>` | 3455 | 17.57672167
`<unnamed wrapped function>` | 46580 | 166.2570884
`init` | 60361 | 26.35722399
`<unnamed wrapped function>` | 78010 | 3.879326344

Also adds new int config functionality.

Fixes #12583
 2022-10-20 00:17:24 +00:00
Peter Hawkins
9ab88071a7 Avoid loading scipy eagerly. 
scipy accounts for around 400ms of the 900ms of JAX's import time. By
loading scipy lazily, we can improve the timing of `import jax` down to
about 500ms.
 2022-10-12 19:51:09 +00:00
Skye Wanderman-Milne
1511439c01 Update version.py and CHANGELOG for jax 0.3.23 release 2022-10-12 11:09:15 -07:00
Skye Wanderman-Milne
e2aa939147 Update Colab TPU driver version 2022-10-11 17:49:10 -07:00
Skye Wanderman-Milne
a6e0e77624 Update version.py, setup.py and CHANGELOG post jax 0.3.22 release 2022-10-11 21:37:53 +00:00
jax authors
af98d8b00f Merge pull request #12746 from jakevdp:fix-changelog 
PiperOrigin-RevId: 480375907
 2022-10-11 09:21:28 -07:00
Jake VanderPlas
05faf0f40d Remove deprecated functionality from jax.test_util 
PiperOrigin-RevId: 480360504
 2022-10-11 08:16:34 -07:00
Jake VanderPlas
943129419c changelog: add missing github commit links 2022-10-11 06:24:38 -07:00
Ran Ran
f3ded0fc1e Address comments for change log 2022-10-05 18:16:49 +00:00
Ran Ran
4870fd3be8 Update message and change log 2022-10-05 04:39:04 +00:00
Skye Wanderman-Milne
2a3460ff8a Update version and CHANGELOG for jax 0.3.21 release 2022-10-03 22:06:19 +00:00
Skye Wanderman-Milne
15e5f38a16 Make persistent compilation cache warn instead of raise an error on cache read/write failures 
Fixes #12582. Setting the env var `JAX_RAISE_PERSISTENT_CACHE_ERRORS=true` will revert to the original behavior of raising exception instead of warning.

Also makes JAX_DUMP_IR_TO work when the persistent cache is enabled.
 2022-09-30 18:38:22 +00:00
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