mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 21:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
12,356 Commits 506 Branches 263 Tags
Commit Graph

12356 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
jax authors
1a7c8831a8 Merge pull request #11589 from skye:workspace 
PiperOrigin-RevId: 462669951
jax-v0.3.15 jaxlib-v0.3.15 jax-v0.3.15-rc 		2022-07-22 11:46:22 -07:00
Skye Wanderman-Milne
26fbeb6e2a Update WORKSPACE and libtpu version for jaxlib 0.3.15, take 3 2022-07-22 11:41:39 -07:00
jax authors
e121e811ab Merge pull request #11536 from sharadmv:colab-debugger 
PiperOrigin-RevId: 462665740
 2022-07-22 11:28:02 -07:00
jax authors
0b6657e471 Merge pull request #11556 from RuffaloLavoisier:tYpO 
PiperOrigin-RevId: 462648717
 2022-07-22 10:13:10 -07:00
Sharad Vikram
4870710891 Enable debugging callbacks with pjit on TPU 
PiperOrigin-RevId: 462527181
 2022-07-21 20:22:14 -07:00
jax authors
8a67734e7b Merge pull request #11579 from sharadmv:fix-effects 
PiperOrigin-RevId: 462478510
 2022-07-21 15:02:46 -07:00
jax authors
7f0b9179f2 Merge pull request #11575 from gnecula:ds_progress 
PiperOrigin-RevId: 462475336
 2022-07-21 14:48:24 -07:00
jax authors
24134ec2a5 Merge pull request #11425 from pschuh:pjit-bugfix 
PiperOrigin-RevId: 462469178
 2022-07-21 14:20:00 -07:00
jax authors
540ee56ff2 Merge pull request #11576 from jakevdp:searchsorted-alt 
PiperOrigin-RevId: 462461853
 2022-07-21 13:47:43 -07:00
Sharad Vikram
d6c172d53e Fix PE not allowing double JIT-ted effectful functions 2022-07-21 11:55:48 -07:00
jax authors
f6c168276b Merge pull request #11578 from jakevdp:wraps-mod 
PiperOrigin-RevId: 462437654
 2022-07-21 11:50:47 -07:00
Jake VanderPlas
9769a0accf DOC: ensure that _wraps() generates correct links to wrapped functions 2022-07-21 11:12:35 -07:00
jax authors
a4e754849e Merge pull request #11543 from nvcastet:fix_multigpu_test 
PiperOrigin-RevId: 462418103
 2022-07-21 10:27:57 -07:00
jax authors
1e05a1cfbc Merge pull request #10816 from mattjj:remove-old-pjit-comment 
PiperOrigin-RevId: 462411602
 2022-07-21 10:01:57 -07:00
George Necula
6c9d2a0b54 [jax2tf] Raise errors for experimental_native_lowering and custom_call 
Raise explicit error when the experimental_native_lowering encounters
a mhlo.custom_call. This would lead to failure when trying to run in TF.
 2022-07-21 19:58:05 +03:00
Jake VanderPlas
10411bfeae jnp.searchsorted: add optional method argument to control implementation 2022-07-21 09:40:18 -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
jax authors
ba7ded4331 Merge pull request #11571 from google:skye-patch-2 
PiperOrigin-RevId: 462267465
 2022-07-20 17:36:20 -07:00
Skye Wanderman-Milne
568cedba8d
Update WORKSPACE for 0.3.15 release, take 2 2022-07-20 17:23:52 -07:00
jax authors
be6db2e619 Merge pull request #10775 from pschuh:mlir-caching 
PiperOrigin-RevId: 462263487
 2022-07-20 17:10:40 -07:00
Parker Schuh
6c4da65af4 Add treedef_is_strict_leaf to fix _prefix_error's semantics. 
Empty nodes like [] and {} have 1 node and 0 leaves. This does not make
them a leaf treedef.

Reproducer:
```
pjit.pjit(lambda x: x, None, (None, {}))((3, {'a': []}))
```
 2022-07-20 17:02:59 -07:00
Kuangyuan Chen
c0ec3b33e6 Introduce jax.experimental.clear_backends to delete all JAX runtime backends. 
In cases like unit tests, users may want to clean up all the backends along with the resources used in the end of the test, and reinitialize them in the next test.

PiperOrigin-RevId: 462239974
 2022-07-20 15:10:27 -07:00
Yash Katariya
d8cbb29d14 OpSharding doesn't have __eq__ defined on it. Don't check sharding equality using opsharding until it does support that. 
PiperOrigin-RevId: 462238497
 2022-07-20 15:03:39 -07:00
Parker Schuh
d8f0099f68 _mlirTransforms merged into _mlirRegisterEverything. 
PiperOrigin-RevId: 462233907
 2022-07-20 14:43:27 -07:00
Yash Katariya
ad67d825fe Add a faster __eq__ check for Mesh. When the id of self and other is the same, there is no need to compare the devices which can be slow when there are 1000s of devices. 
PiperOrigin-RevId: 462230016
 2022-07-20 14:25:41 -07:00
Yash Katariya
026636951a Add lru_cache and use it instead of util.cache() in places where tracing user code is not required. 
Co-authored-by: Matthew Johnson <mattjj@google.com>
PiperOrigin-RevId: 462212010
 2022-07-20 13:05:20 -07:00
jax authors
ffe67c1042 Merge pull request #11563 from jakevdp:upstream-ci 
PiperOrigin-RevId: 462200130
 2022-07-20 12:09:08 -07:00
jax authors
5370f59036 Merge pull request #11566 from jakevdp:missing-f 
PiperOrigin-RevId: 462187015
 2022-07-20 11:16:22 -07:00
jax authors
d0162cd37e Merge pull request #11533 from ROCmSoftwarePlatform:rocm_disable_lobpcg_test 
PiperOrigin-RevId: 462182051
 2022-07-20 10:56:54 -07:00
Jake VanderPlas
114b03670c Add missing f-string marker 2022-07-20 10:48:07 -07:00
Jake VanderPlas
993196c451 CI: make parse_logs more robust to errors 2022-07-20 10:33:32 -07:00
Adam Paszke
117da44712 Internal change 
PiperOrigin-RevId: 462110048
 2022-07-20 04:31:21 -07:00
RuffaloLavoisier
9f770425ac Correct spelling on word 2022-07-20 18:57:12 +09:00
jax authors
7f1813c5e3 Merge pull request #11539 from gnecula:ds_reshape 
PiperOrigin-RevId: 462061742
 2022-07-19 23:13:03 -07:00
jax authors
9ee6cacdc8 Merge pull request #11540 from gnecula:ds_check_flag 
PiperOrigin-RevId: 462061356
 2022-07-19 23:07:14 -07:00
jax authors
02c3d6bf2c Merge pull request #11553 from mattjj:readme-windows-update 
PiperOrigin-RevId: 462022337
 2022-07-19 18:19:37 -07:00
jax authors
7e5bc2977b Merge pull request #11552 from mattjj:mhlo-bint-progress 
PiperOrigin-RevId: 462015062
 2022-07-19 17:32:00 -07:00
jax authors
5a9b8490eb Merge pull request #11554 from google:skye-patch-1 
PiperOrigin-RevId: 462013784
 2022-07-19 17:24:56 -07:00
Matthew Johnson
3ab8637a73 Update README discussion of Windows support: 
* right next to the pip installation instructions, mention they don't work for Windows;
* add a link to #5795 for an unofficial discussion of Windows native support
 2022-07-19 17:12:40 -07:00
Skye Wanderman-Milne
186a4f83e3
Update libtpu version for 0.3.15 release 2022-07-19 17:11:43 -07:00
jax authors
18541e2efa Merge pull request #11542 from mattjj:remove-resnet50-example 
PiperOrigin-RevId: 462008290
 2022-07-19 16:56:56 -07:00
Matthew Johnson
7cb5c2447e [dynamic-shapes] fix minor bint bugs 
Co-authored-by: Eugene Burmako <burmako@google.com>
 2022-07-19 16:38:40 -07:00
jax authors
6ea9e4d4dd Merge pull request #11546 from jakevdp:fix-scipy-sym-pos 
PiperOrigin-RevId: 461976166
 2022-07-19 14:26:51 -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
Yash Katariya
9f914a93d6 Replace input sharding_specs with in_shardings in InputsHandler 
PiperOrigin-RevId: 461963206
 2022-07-19 13:43:28 -07:00
jax authors
226cef08bf Merge pull request #11544 from jakevdp:conv-general-pad 
PiperOrigin-RevId: 461937001
 2022-07-19 11:46:08 -07:00
Jake VanderPlas
489596c0e2 lax.conv_general_dilated: validate negative paddings 2022-07-19 11:15:18 -07:00
George Necula
2106d65561 [dynamic-shapes] Add check that --jax_dynamic_shapes is set when using abstracted_axes. 
abstracted_axes has no effect without the --jax_dynamic_shapes. Make this and
explicit error.
 2022-07-19 19:48:45 +02:00
jax authors
ac731bb4cc Merge pull request #11518 from ROCmSoftwarePlatform:rocm_disable_dlpack_unit_tests 
PiperOrigin-RevId: 461904181
 2022-07-19 09:48:36 -07:00
jax authors
97fe2262d3 Merge pull request #11526 from atgctg:patch-2 
PiperOrigin-RevId: 461903614
 2022-07-19 09:42:30 -07:00