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

41 Commits

Author SHA1 Message Date
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
Peter Hawkins
48bbdbc890 Change jax.core.DropVar to be a non-singleton. 
Previously jax.core.DropVar was a singleton value (jax.core.dropvar) whose type was always jax.core.AbstractUnit. However, this type is misleading: a DropVar is an equation output, and typically we would expect it to have an array type. In particular, the unit type confuses new-style translation rules that expect to use the output aval on an equation as part of the lowering logic.

Instead, change DropVar to be a non-singleton subclass of Var instead with a flexible choice of aval.

PiperOrigin-RevId: 404071001
 2021-10-18 15:02:54 -07:00
Jake VanderPlas
f2bbd51cc2 [sparse] respect weak types in sparsify transform 2021-10-18 13:35:20 -07:00
Jake VanderPlas
c2dd90e3a0 [sparse] Factor BCOO-related routines into a separate submodule 2021-10-06 08:06:18 -07:00
Jake VanderPlas
3a440d665f [sparse] add sparsify support for sparse-sparse matmul 2021-10-05 16:45:48 -07:00
Peter Hawkins
e869e5e0f8 Move contents of jax.api_util to jax._src.api_util and add a forwarding shim. 
One of many changes to codify the set of exported symbols in the jax.* namespace.

PiperOrigin-RevId: 395484706
 2021-09-08 09:00:56 -07:00
jax authors
b9cc31e35d Merge pull request #7852 from google:sparse-jaxpr-consts 
PiperOrigin-RevId: 395421332
 2021-09-08 01:27:01 -07:00
Roy Frostig
8bb8bf1081 avoid constvar conversion when closing a sparse jaxpr 2021-09-07 22:02:21 -07:00
Roy Frostig
bf44398790 handle dropped output values in the sparse interpreter 2021-09-07 18:50:13 -07:00
Jake VanderPlas
82a7b7ee4d DOC: add documentation of jax.experimental.sparse 2021-09-02 17:08:10 -07:00
Jake VanderPlas
c5fed9c3b5 [sparse] Change BCOO index order 2021-09-01 13:48:55 -07:00
Jake VanderPlas
4f9310088d [sparse] handle pytree inputs in sparsify transform 2021-08-10 10:31:16 -07:00
Jake VanderPlas
0b7c0daee2 [sparse] bug: thread through params in sparsify 2021-08-09 12:14:37 -07:00
Jake VanderPlas
25b3737e81 [sparse] correctly handle units in sparsify argspecs 2021-08-09 09:15:08 -07:00
Jake VanderPlas
f76108ba0e [sparse] add sparsify rule for lax.cond 2021-08-06 13:32:23 -07:00
Jake VanderPlas
1d359f8c61 [sparse]: add sparse rule for scan/fori_loop 2021-08-05 15:19:43 -07:00
Jake VanderPlas
1eb3b5f8d6 [sparse] support sparse arguments in xla_call 2021-07-13 15:23:14 -07:00
Jake VanderPlas
9af8676341 [sparse] support dense xla_call within sparsify jaxpr interpreter 2021-07-13 13:31:21 -07:00
Jake VanderPlas
5db97e0bf9 [sparse] add sparse transform rule for lax.while_p 2021-07-12 16:53:24 -07:00
Jake VanderPlas
087da553cd [sparse]: add support for rdot_general in sparsify transform 2021-07-09 06:00:05 -07:00
Jake VanderPlas
76f9e6f016 [sparse] globally change nnz->nse 2021-06-30 17:46:02 -07:00
Jake VanderPlas
5ed9471b9a flake: fix unused import 2021-06-28 11:40:23 -07:00
jax authors
cb8582c63d Merge pull request #6929 from jakevdp:sparsify 
PiperOrigin-RevId: 381897788
 2021-06-28 10:43:07 -07:00
Jake VanderPlas
0401d2be57 Add experimental sparsify transform 
Co-authored-by: Roy Frostig <frostig@google.com>
 2021-06-25 10:45:16 -07:00