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

13773 Commits

This Branch
This Branch
All Branches
Author SHA1 Message Date
Peter Hawkins
aa658bde6f Disable asan/tsan for a number of slow tests. 
PiperOrigin-RevId: 488356786
 2022-11-14 07:12:16 -08:00
Peter Hawkins
40e81c3a86 Revert: Use pinv to compute lstsq. 
The current implementation of lstsq is equivalent to pinv(A) @ b, with a different order of matrix multiplications. If we write it that way we benefit from a more stable derivative that does not require differentiating through the singular value decomposition.

This PR appears to have caused numerical problems in downstream tests.

PiperOrigin-RevId: 487942754
 2022-11-11 16:24:05 -08:00
jax authors
d4cc5882d2 Merge pull request #13215 from ROCmSoftwarePlatform:sytrd 
PiperOrigin-RevId: 487940256
 2022-11-11 16:11:27 -08:00
Sharad Vikram
e15619ceab Convert string axis name into tuple of strings in Mesh constructor 
PiperOrigin-RevId: 487930412
 2022-11-11 15:27:51 -08:00
Yash Katariya
6897d37562 Add docstrings for jax.Array APIs make_array_from_callback and make_array_from_single_device_arrays. 
PiperOrigin-RevId: 487929688
 2022-11-11 15:21:10 -08:00
Rahul Batra
31d8f62826 Sytrd solver and SytrdDescriptor should NOT be CUDA only 2022-11-11 22:41:51 +00:00
jax authors
19d76a7818 Merge pull request #13212 from sharadmv:fix-changelog 
PiperOrigin-RevId: 487915320
 2022-11-11 14:18:19 -08:00
Sharad Vikram
4bdfdd7363 Update changelog w/ info about deleting jax_experimental_name_stack 2022-11-11 14:02:30 -08:00
Peter Hawkins
7c3fb81310 Use pinv to compute lstsq. 
The current implementation of lstsq is equivalent to pinv(A) @ b, with a different order of matrix multiplications. If we write it that way we benefit from a more stable derivative that does not require differentiating through the singular value decomposition.

PiperOrigin-RevId: 487903227
 2022-11-11 13:28:48 -08:00
Peter Hawkins
047974dd0c Be more economical when computing the JVP of the SVD of non-square matrices. 
(Note this isn't a regression from #13147: the previous change did not alter the order of operations.)

PiperOrigin-RevId: 487896154
 2022-11-11 12:55:53 -08:00
Peter Hawkins
c9ebf60f4e Compute the JVP of jnp.linalg.pinv more economically for non-square matrices. 
The order of the matrix products matters.

PiperOrigin-RevId: 487879202
 2022-11-11 11:45:02 -08:00
jax authors
995736119e Merge pull request #13198 from patrick-kidger:prng-isinstance 
PiperOrigin-RevId: 487859092
 2022-11-11 10:31:24 -08:00
jax authors
c359c7976b Merge pull request #13196 from jakevdp:simpler-sparsify 
PiperOrigin-RevId: 487837115
 2022-11-11 09:01:55 -08:00
Patrick Kidger
d2afa84a6e PRNGKeyArray is now a virtual subclass of ndarray 2022-11-11 08:04:38 -08:00
Felix Chern
10e6fe8cde Correct norm in ann.py doc. 
PiperOrigin-RevId: 487814084
 2022-11-11 07:08:54 -08:00
jax authors
ce85106578 Merge pull request #13193 from tlu7:bcsr-fromdense-batching 
PiperOrigin-RevId: 487810600
 2022-11-11 06:51:30 -08:00
jax authors
71f92a7cd5 Merge pull request #13182 from canyon289:docs_update 
PiperOrigin-RevId: 487808462
 2022-11-11 06:44:29 -08:00
jax authors
bdf3bd5472 Merge pull request #13155 from jakevdp:bcoo-gather 
PiperOrigin-RevId: 487808456
 2022-11-11 06:37:30 -08:00
Jake VanderPlas
ea25b79b87 [sparse] streamline sparse rules for standard primitives 2022-11-11 04:50:33 -08:00
Jake VanderPlas
90dc008340 [sparse] add bcoo_gather & support for sparse indexing 2022-11-11 04:25:14 -08:00
Peter Hawkins
a13541441b Reenable a TPU test now that the compiler bug is fixed. 
PiperOrigin-RevId: 487705048
 2022-11-10 19:38:01 -08:00
Tianjian Lu
332fced0cc sparse] BCSR batching rule. 
[Co-authored-by: Jake Vanderplas: <vanderplas@google.com>
 2022-11-10 19:33:32 -08:00
jax authors
dc0d7ba368 Merge pull request #13202 from google:yashk2810-patch-18 
PiperOrigin-RevId: 487701933
 2022-11-10 19:15:29 -08:00
Yash Katariya
f0c0689a8a
Remove internal information 2022-11-10 19:09:19 -08:00
Parker Schuh
4a3b7f16ff Change pickling for jax.sharding to not serialize device ids. 
PiperOrigin-RevId: 487700467
 2022-11-10 19:05:02 -08:00
jax authors
f9d7a6ae20 Merge pull request #13197 from google:yashk2810-patch-17 
PiperOrigin-RevId: 487687006
 2022-11-10 17:56:29 -08:00
Yash Katariya
73935a5bd1
Update jax_array_migration.md 2022-11-10 17:23:16 -08:00
jax authors
c318f771cb Merge pull request #13185 from tlu7:bcsr-from-scipy 
PiperOrigin-RevId: 487680265
 2022-11-10 17:18:44 -08:00
Yash Katariya
aa66b939f9
Fix the docs build 2022-11-10 17:08:57 -08:00
Yash Katariya
b49a1bda15 Add jax.Array migration doc to OSS 
PiperOrigin-RevId: 487673643
 2022-11-10 16:46:30 -08:00
Tianjian Lu
311fb24ff9 [sparse] Add BCSR from_scipy_sparse. 
Co-authored-by: Jake Vanderplas <vanderplas@google.com>
 2022-11-10 16:44:59 -08:00
Peter Hawkins
352b042fe9 Add a GPU implementation of symmetric (Hermitian) tridiagonal reduction. 
Change the contract of lax.linalg.tridiagonal to return the d and e vectors as well. Since we only just added this function and have never released JAX with it we can make this change without breaking compatibility.

Also fix wrong dtypes for d and e values in the CPU lapack sytrd wrapper.

PiperOrigin-RevId: 487621469
 2022-11-10 13:16:21 -08:00
Sharad Vikram
74b136e62c Delete jax_experimental_name_stack flag 
PiperOrigin-RevId: 487601864
 2022-11-10 11:59:50 -08:00
jax authors
0ebb6b4215 Merge pull request #13180 from jakevdp:bcoo-slice 
PiperOrigin-RevId: 487568853
 2022-11-10 10:04:35 -08:00
Yash Katariya
cc41ee85c4 Mark scipy_signal_test and sparse_test optonly because it times out under debug mode. 
PiperOrigin-RevId: 487533356
 2022-11-10 07:38:58 -08:00
Yash Katariya
71360edf90 Bump the shard count for TPU to avoid timeouts 
PiperOrigin-RevId: 487421018
 2022-11-09 20:32:12 -08:00
Peter Hawkins
e42e52d4aa Rename test flag --num_generated_cases to --jax_num_generated_cases. 
parse_flags_with_absl() only parses flags that start with --jax_. Other flags are only parsed when absl.app's main function runs. But that's too late for test cases: test cases need to have the number of generated cases chosen at module initialization time. Hence the --num_generated_cases flag wasn't doing anything. Oops. By renaming it it works once again.

It might make sense to stop using flags for the number of generated cases and only use environment variables. We defer that to a future change.

Fix many test cases that were shown to be broken with a larger number of test cases enabled.

PiperOrigin-RevId: 487406670
 2022-11-09 18:58:05 -08:00
jax authors
b36afc5b0d Merge pull request #13177 from jakevdp:bcoo-dynamic-slice 
PiperOrigin-RevId: 487390430
 2022-11-09 17:29:30 -08:00
Sharad Vikram
3731e446c0 Set default layout for Python callback 
PiperOrigin-RevId: 487388682
 2022-11-09 17:18:49 -08:00
Yash Katariya
f9bbd585b9 Improve the error message when @pjit (with no {in_axis|out_axis}_resources is used without jax.Array enabled. 
PiperOrigin-RevId: 487380328
 2022-11-09 16:38:00 -08:00
Ravin Kumar
cb3762e5dd Consolidate links in JAX documentation 
Move notes down
 2022-11-09 15:44:52 -08:00
Jake VanderPlas
4c4f2a3ad2 [sparse] support strides in bcoo_slice 2022-11-09 15:03:21 -08:00
John QiangZhang
b41c594a19 Internal change 
PiperOrigin-RevId: 487351591
 2022-11-09 14:41:39 -08:00
Jake VanderPlas
46d9cac122 [sparse] bcoo_dynamic_slice: remove unnecessary padding from output 2022-11-09 13:56:18 -08:00
jax authors
fa0217bdd3 Merge pull request #13175 from jakevdp:bcoo-transpose 
PiperOrigin-RevId: 487333413
 2022-11-09 13:34:55 -08:00
Jake VanderPlas
0c3e330148 [sparse] fix shape bug in bcoo_transpose 2022-11-09 12:53:13 -08:00
Felix Chern
8ac7422e26 [JAX] Disables large k test cases in ann_test. 
Will investigate probability properties for the corner cases in the future.

PiperOrigin-RevId: 487302143
 2022-11-09 11:32:47 -08:00
jax authors
63e3152764 Merge pull request #13160 from jakevdp:bcoo-squeeze 
PiperOrigin-RevId: 487280563
 2022-11-09 10:18:22 -08:00
jax authors
f697b8e087 Merge pull request #13166 from LenaMartens:checking-keys 
PiperOrigin-RevId: 487267267
 2022-11-09 09:31:39 -08:00
lenamartens
053b8b5bcd Checkify: fix nan_checks+PRNGKeys - a PRNGKey is never NaN! 
Add a guard to the nan_error_rule to not call jnp.isnan on keys.
 2022-11-09 17:08:21 +00:00