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
|
0b5443c6e8
|
Clean up: remove unused helper functions
|
2023-02-01 09:55:58 -08:00 |
|
Jake VanderPlas
|
c0c347bbe7
|
[sparse] add nfold_vmap utility
|
2023-01-12 10:09:58 -08:00 |
|
Tianjian Lu
|
49086e8455
|
[sparse] Use SparseInfo for both bcoo and bcsr format.
PiperOrigin-RevId: 495968517
|
2022-12-16 15:00:38 -08:00 |
|
Jake VanderPlas
|
c5a6ec1705
|
[sparse] fix count_stored_elements for empty batches
|
2022-12-14 16:49:05 -08:00 |
|
Jake VanderPlas
|
e9cc523873
|
[sparse] validate BCOO on instantiation
|
2022-12-12 11:39:06 -08:00 |
|
Tianjian Lu
|
942aa7a907
|
[sparse] Move _dot_general_validated_shape to sparse util.
PiperOrigin-RevId: 494031113
|
2022-12-08 16:54:43 -08:00 |
|
Jake VanderPlas
|
c376836721
|
[typing] annotate jax.experimental.sparse
|
2022-12-05 14:05:25 -08:00 |
|
Tianjian Lu
|
58a2abe1b5
|
[sparse] Move broadcasting_vmap to sparse util.
PiperOrigin-RevId: 478566197
|
2022-10-03 11:56:28 -07:00 |
|
Tianjian Lu
|
67b7ae259f
|
[sparse] Move _bcoo_nse to sparse util.
PiperOrigin-RevId: 476263483
|
2022-09-22 20:22:06 -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
|
c1549a0a16
|
[sparse] make sparse objects compatible with jax.jit.lower()
|
2022-07-15 09:58:31 -07:00 |
|
Jake VanderPlas
|
bf4c3b64af
|
[sparse] add bcoo_update_layout utility
|
2022-05-05 14:53:38 -07: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 |
|