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
16,052 Commits 506 Branches 263 Tags
Commit Graph

1343 Commits

Author SHA1 Message Date
Roy Frostig
ce840a9cd8 JEP: jax.extend, a module for extensions 2023-05-05 13:50:22 -07:00
jax authors
5d143e6eea Merge pull request #15818 from froystig:random-bits-direct 
PiperOrigin-RevId: 529090390
 2023-05-03 07:56:17 -07:00
Roy Frostig
ea3389205f add jax.random.bits 2023-05-03 06:10:05 -07:00
David Pizzuto
6948d32d15 contributing: Switch repo URL to HTTPS for consistency with other github URLs. 2023-05-01 10:03:39 -07:00
Jake VanderPlas
e059e3b52f DOC: document jax.experimental.sparse.linalg 2023-04-28 14:18:50 -07:00
Jake VanderPlas
8dc06ed2ce Document jax.lax.with_sharding_constraint 2023-04-26 10:19:04 -07:00
jax authors
70ebdb0502 Jax101 - Clarify that the compiled code is executed on first call 
The current wording implies that the first time a jitted function
is called, the computation happens in Python. It's actually only
the tracing that happens in Python, and the compiled code is run
during the first call. The distinction is important e.g., to
understand why it might make sense to jit a function that's only
called once.

PiperOrigin-RevId: 526906176
 2023-04-25 02:41:12 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
jax authors
975e76ef76 Merge pull request #15664 from skye:tpu_install 
PiperOrigin-RevId: 525605301
 2023-04-19 18:18:32 -07:00
jax authors
1de4d14da8 Merge pull request #15656 from laqua-stack:add-special-gamma-fcn 
PiperOrigin-RevId: 525566749
 2023-04-19 15:28:36 -07:00
Jake VanderPlas
a083ba7853 DOC: explicitly mention io_callback in FAQ 2023-04-19 12:30:53 -07:00
Skye Wanderman-Milne
b917a31f56 Update TPU install on main docs page 2023-04-19 17:52:16 +00:00
laqua-stack
d742733bea feat (scipy.special): Add a xla version of scipy.special.gamma function 
- Add gamma fcn api in scipy.special
- Add tests for this purpose
- Add function to the docs

Currently, there is no implementation of the gamma function in jax
but there is one in scipy.special. This breaks some higher level
jit-compilation like in the blackjax backend for pymc. This commit
adds the missing gamma function.

Resolves: #15409
 2023-04-18 21:10:22 +02:00
jax authors
cabf8b7302 Merge pull request #15636 from Vaishaal:idct 
PiperOrigin-RevId: 524936117
 2023-04-17 13:35:34 -07:00
Vaishaal Shankar
add15aca25 implement idct and idctn + add function to scipy.rst 2023-04-17 12:12:51 -07:00
jax authors
9c92ec5d90 Merge pull request #15585 from SauravMaheshkar:main 
PiperOrigin-RevId: 524896086
 2023-04-17 11:08:46 -07:00
Matthew Johnson
690071f1de fix custom_jvp docs typo 2023-04-15 14:51:33 -07:00
Saurav Maheshkar
cfd8762d4a feat: move configurations to pyproject 2023-04-15 02:39:39 +01:00
Jake VanderPlas
055edf4a08 DOC: add docstrings for callback functions 2023-04-12 07:33:09 -07:00
jax authors
5456152f0b Fix typo: "__W__ - spatial height" -> "__W__ - spatial width" 
PiperOrigin-RevId: 523365127
 2023-04-11 04:37:26 -07:00
Matthew Johnson
a7f5e07549 update prng docs to mention jax_threefry_partitionable 
fixes #15484
 2023-04-07 22:55:15 -07:00
Matthew Johnson
057d408448 add docs for jax.clear_caches 
Co-authored-by: Roy Frostig <frostig@google.com>
 2023-04-07 14:42:31 -07:00
jax authors
5f73b9e029 Merge pull request #15440 from jakevdp:callback-doc-2 
PiperOrigin-RevId: 522678231
 2023-04-07 14:01:23 -07:00
jax authors
2ebb178c35 Merge pull request #15224 from jecampagne:fftconvolve2dr 
PiperOrigin-RevId: 522671725
 2023-04-07 13:29:50 -07:00
Dan Kondratyuk
920e761c40 Add info about iterating over collections with non-deterministic ordering. 
PiperOrigin-RevId: 522611241
 2023-04-07 09:12:49 -07:00
Jean-Eric Campagne
4beee13ba0 Add implementation of jax.scipy.fftconvolve 2023-04-07 17:19:08 +02:00
Jake VanderPlas
4473ebc9fc Add documentation for callback functions 2023-04-07 07:05:46 -07:00
jax authors
8a6c929678 Merge pull request #15289 from cgarciae:add-missing-api-references 
PiperOrigin-RevId: 521617419
 2023-04-03 18:23:04 -07:00
jax authors
ffb8352848 Merge pull request #15342 from jakevdp:doc-requirements 
PiperOrigin-RevId: 520955387
 2023-03-31 10:27:30 -07:00
Jake VanderPlas
9ec3ad1ce7 DOC: pin newest sphinx-book-theme 2023-03-31 08:42:34 -07:00
Jake VanderPlas
749dc1b95e Remove deprecated function jnp.msort 2023-03-31 08:24:36 -07:00
Jake VanderPlas
ec63d699e9 DOC: fix headings in user_guides 2023-03-30 14:39:25 -07:00
jax authors
dedfc8df75 Merge pull request #15282 from JiaYaobo:geom_random 
PiperOrigin-RevId: 520635974
 2023-03-30 07:45:19 -07:00
Jake VanderPlas
8562a8d7bc DOC: pin pydata-sphinx-theme to fix incompatibility 2023-03-29 13:55:32 -07:00
jiayaobo
924894e85c add geometric random gen 
add geom random

add geom random

add geom random

add geom random
 2023-03-30 02:08:04 +08:00
Yash Katariya
fbc05ee5ac Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago 
PiperOrigin-RevId: 520356179
 2023-03-29 09:23:22 -07:00
Cristian Garcia
1e7f9d9e78 add key_path functions to tree_util api reference 2023-03-29 15:17:13 +00:00
Peter Hawkins
d9b0f3cd6f Recommend --local_test_jobs in bazel test command line on GPU. 2023-03-29 09:28:53 -04:00
Skye Wanderman-Milne
473d1c3685 Turn on PJRT C API by default. 
I forgot that the default setting is actually in jaxlib:
fbe9a80fdb/xla/python/xla_client.py (L135)

To be able to make this change as a jax-only release, I manually set
the env var on Cloud TPU if it isn't already set.
 2023-03-28 15:28:13 -07:00
jax authors
af4d4943a7 Merge pull request #8633 from shawwn:2021-11-19/autodidax-fix-jaxpr-subcomp-return-type 
PiperOrigin-RevId: 519745476
 2023-03-27 09:52:20 -07:00
Ravin Kumar
8c2549519b
Update user_guides.rst 
Fix minor typo
 2023-03-26 17:21:35 -07:00
Peter Hawkins
6ed66ada0f Delete remote TPU support. 
TPU VMs are the only supported way to use TPUs as of JAX 0.4.0.

PiperOrigin-RevId: 519211267
 2023-03-24 12:33:33 -07:00
jax authors
383cf41848 Merge pull request #14937 from b0nce:fix-stats 
PiperOrigin-RevId: 518888600
 2023-03-23 10:01:54 -07:00
jax authors
00e6c73b68 Merge pull request #15114 from JiaYaobo:add_wald_random 
PiperOrigin-RevId: 518592428
 2023-03-22 09:37:44 -07:00
jiayaobo
f7a14d65d2 add wald random generator 
add wald to random.py
 2023-03-22 11:06:59 +08:00
Jake VanderPlas
4a9ed3eaa8 Document ShapeDtypeStruct 2023-03-21 13:53:20 -07:00
Misha
83b3f5b759 Fix loc and scale parameters in scipy.logistic. Add CDF and SF for several distributions. 2023-03-21 00:16:13 +01:00
jax authors
fa9d9ae05f Merge pull request #14900 from JiaYaobo:add_rayleigh_random 
PiperOrigin-RevId: 518015562
 2023-03-20 10:51:49 -07:00
Mark Sandler
bab1098866 Fixes broken examples, and (invalid) comment for PartitionSpec 
PiperOrigin-RevId: 517531823
 2023-03-17 16:09:45 -07:00
jax authors
c25ea3f0f2 Merge pull request #15064 from jakevdp:sharp-bits-indexing 
PiperOrigin-RevId: 517498861
 2023-03-17 13:50:14 -07:00