mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-25 19:16:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
15,071 Commits 508 Branches 263 Tags
Commit Graph

73 Commits

Author SHA1 Message Date
Jake VanderPlas
f3e5024787 [sparse] implement bcsr_concatenate 2023-02-15 14:10:56 -08:00
Jake VanderPlas
d688b6d6f3 [sparse] implement bcsr_broadcast_in_dim 2023-02-15 10:26:10 -08:00
Roy Frostig
cb8dcce2fe migrate more internal dependencies from jax.core to jax._src.core 
PiperOrigin-RevId: 509736368
 2023-02-14 23:01:11 -08:00
Jake VanderPlas
de8a77a3eb [sparse] implement BCSR.__matmul__ 2023-02-10 16:11:57 -08:00
Jake VanderPlas
597c20173f [sparse] support BCSR in sparsify transform 2023-02-06 11:01:57 -08:00
Jake VanderPlas
428713e88e [sparse] support all unbatched 1D convolutions 2023-02-03 15:59:42 -08:00
Jake VanderPlas
038798ed25 [sparse] add support for simple 1D convolutions 2023-02-01 18:53:49 -08:00
Jake VanderPlas
4fa80b44cd [sparse] implement sparse rule for lax.rev 2023-02-01 15:43:47 -08:00
Jake VanderPlas
05b7189873 [sparse] add a few missing methods to BCOO 2023-01-20 10:26:57 -08:00
Yash Katariya
c3bb26050c Add pjit rule to sparse_rules to support pjit. This is done to merge the jit and pjit API. 
PiperOrigin-RevId: 499311841
 2023-01-03 14:13:19 -08:00
Jake VanderPlas
4a6bbde409 Move jax.linear_util to jax._src.linear_util 2022-12-20 14:49:27 -08:00
Jake VanderPlas
c376836721 [typing] annotate jax.experimental.sparse 2022-12-05 14:05:25 -08:00
Jake VanderPlas
e673f1fd44 [sparse] avoid re-indexing for linear unary ops 2022-11-17 16:31:46 -08:00
Jake VanderPlas
c8ee485de8 [sparse] implement iter() of BCOO 2022-11-15 13:17:06 -08:00
Yash Katariya
a419e1917a Use jax.Array by default for doctests 
PiperOrigin-RevId: 488719467
 2022-11-15 11:52:22 -08:00
jax authors
c359c7976b Merge pull request #13196 from jakevdp:simpler-sparsify 
PiperOrigin-RevId: 487837115
 2022-11-11 09:01:55 -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
Jake VanderPlas
4255697610 [sparse] add bcoo_squeeze function 2022-11-08 18:16:20 -08:00
Jake VanderPlas
7c0d0e67c8 [sparse] add support for BCOO.astype method 2022-11-08 13:30:22 -08:00
Jake VanderPlas
6ed9b14d8d [sparse] simplify BCOO indexing implementation 2022-11-02 10:13:35 -07:00
Jake VanderPlas
5d15757741 [typing] annotate jax._src.util.safe_map 2022-10-20 10:15:04 -07:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". 
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
 2022-09-22 12:27:19 -07:00
Jake VanderPlas
47b9f216bc [sparse] add sparse support for dynamic_slice 2022-09-01 13:42:02 -07:00
Jake VanderPlas
269e75273a [sparse] add support for bcoo equivalent of lax.slice 2022-08-25 09:19:44 -07:00
Jake VanderPlas
d78b96937c [sparse] add trivial sparse rule for lax.copy_p 2022-08-19 16:41:18 -07:00
Jake VanderPlas
2b4f72b6f4 [sparse] fix unary operations in presence of duplicate indices 2022-07-07 13:49:50 -07:00
Tianjian Lu
cc4f2ade63 [sparse] Track unique_indices in sparse transform. 
PiperOrigin-RevId: 452848200
 2022-06-03 14:51:26 -07:00
Tianjian Lu
b6d4c59e64 [sparse] Trace BCOO indices_sorted in sparsifying zero_preserving_unary_ops. 
PiperOrigin-RevId: 451081409
 2022-05-25 21:07:29 -07:00
Tianjian Lu
1656b33ca9 [sparse] Track indices_sorted in sparsify transform. 
PiperOrigin-RevId: 449833669
 2022-05-19 14:29:14 -07:00
Jake VanderPlas
1419c7a0d5 [sparse] raise better error for subtraction between sparse and dense 2022-05-17 09:50:44 -07:00
jax authors
03c80209d2 Merge pull request #10532 from mattjj:remove-units-final 
PiperOrigin-RevId: 446519086
 2022-05-04 12:09:51 -07:00
Matthew Johnson
9cd55a2bbd [remove-units] remove units 2022-05-04 10:58:56 -07:00
Jake VanderPlas
1a9a796a0c [sparse] implement sparse rule for lax.reshape_p 2022-05-02 09:11:55 -07:00
Jake VanderPlas
2d9af38a2c [sparse] implement sparse rule for lax.concatenate_p 2022-04-28 10:59:11 -07:00
Jake VanderPlas
c37c1e683e [sparse] improve error messages for unimplemented primitives 2022-04-25 16:22:19 -07:00
Jake VanderPlas
01e4fa8a78 [sparse] consolidate flavors of bcoo_dot_general 2022-04-07 11:28:12 -07:00
Jake VanderPlas
aa0da8e8e7 [sparse] make bcoo_spdot_general return a BCOO array, not raw buffers 2022-04-06 15:40:26 -07:00
Tianjian Lu
e5d7f65b6a [sparse] Change call signature of bcoo primitive wrappers. 
PiperOrigin-RevId: 437923482
 2022-03-28 20:42:52 -07:00
Jake VanderPlas
8c6e001e45 [sparse] refactor internal implementation of sparsify transform 2022-03-07 12:48:03 -08:00
Jake VanderPlas
2c20d82776 [sparse] generalize metadata argument in BCOO primitives 2022-01-07 14:07:55 -08:00
Jake VanderPlas
b58ac44228 [sparse] add sparse rule for lax.sub_p 2021-12-13 16:46:47 -08:00
jax authors
7ceb983bc8 Merge pull request #8795 from jakevdp:bcoo-sparse-mul 
PiperOrigin-RevId: 414734890
 2021-12-07 08:38:31 -08:00
Jake VanderPlas
242463f5c5 [sparse] add sparse rule for broadcast_in_dim 2021-12-06 17:25:28 -08:00
Jake VanderPlas
74a512a9cf [sparse] initial implementation of sparse-sparse multiply 2021-12-06 09:02:34 -08:00
Jake VanderPlas
da319cf302 [sparse] refactor jax.experimental.sparse 
Why? Better organization, and to avoid issues with circular imports.

PiperOrigin-RevId: 413679493
 2021-12-02 08:02:15 -08:00
Jake VanderPlas
b00061e038 [sparse]: add tracer-based implementation of sparsify 
Co-authored by: Matthew Johnson <mattjj@google.com>
 2021-11-04 13:02:46 -07:00
Jake VanderPlas
4d4dae184f [sparse] implement broadcasted sparse-dense multiplication 2021-10-26 14:45:21 -07:00
Jake VanderPlas
2259e2b0a8 [sparse] add todense() primitive for use in sparsify transform 2021-10-26 13:52:48 -07:00
jax authors
fa5d4a2a07 Merge pull request #8268 from jakevdp:sparsify-weak-types 
PiperOrigin-RevId: 404097482
 2021-10-18 16:49:41 -07:00