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

2123 Commits

Author SHA1 Message Date
Jake VanderPlas
f5ca46f5ec Sharp bits: add note on subnormal flush-to-zero 2025-02-24 18:52:38 -08:00
Peter Hawkins
85add667ed Change documentation to recommend libtpu from pypi instead of GCS. 2025-02-24 17:50:29 -05:00
George Necula
1be801bac8 [better_errors] Cleanup use of DebugInfo.arg_names and result_paths 
Previously, we represented a missing arg name with `None`,
and a missing result path with the empty string. We now
adopt the same convention for arg names and use empty strings.
This simplifies the typing, and prevents the string "None" from
appearing in error messages.

I changed how we encode the result paths. Previously for a
function that returns a single array the path was the empty
string (the same as for an unknown path). And for a function
that returns a pair of arrays it was `([0], [1])`. Now we
add the "result" prefix: `("result",)` for a function returning a
single array and `(result[0], result[1])` for a function returning
a pair of arrays.

Finally, in debug_info_test, I removed the `check_tracer_arg_name`
so that all spied tracers are printed with the argument name they
depend on.
 2025-02-23 08:27:56 +02:00
Jake VanderPlas
380fb0a5a0 CI: skip execution of convolutions.ipynb 
This is failing on readthedocs with a dead kernel error
 2025-02-20 13:14:59 -08:00
rajasekharporeddy
180be99798 Fix typos 2025-02-18 17:01:29 +05:30
Dan Foreman-Mackey
c6c38fb852 Reorder top-level functions in lax.linalg, and add/expand docstrings. 
PiperOrigin-RevId: 726603731
 2025-02-13 12:57:55 -08:00
Dougal
9145366f6f Part 1 of a new autodidax based on "stackless" 2025-02-12 15:23:06 -05:00
Roy Frostig
8720a9c0cd docstrings and API reference doc listing for the traced AOT stage 2025-02-11 22:30:50 -08:00
jax authors
914adaf60c Merge pull request #26476 from froystig:aot-doc-traced 
PiperOrigin-RevId: 725902103
 2025-02-11 22:01:21 -08:00
Roy Frostig
af381a73a3 update AOT walkthrough to cover explicit tracing stage 2025-02-11 21:26:05 -08:00
Jake VanderPlas
e389b707ba Add public APIs for jax.lax monoidal reductions 2025-02-11 16:00:03 -08:00
Jiyoun (Jen) Ha
e6f38b7fb3 [pallas] fix typo 
PiperOrigin-RevId: 725244824
 2025-02-10 09:32:07 -08:00
Dan Foreman-Mackey
21281f377e Remove the tutorial about legacy custom calls. 2025-02-06 21:04:30 -05:00
Jake VanderPlas
e0ee415677 DOC: update numpy/scipy versions in deprecation doc 2025-02-06 15:18:11 -08:00
Jake VanderPlas
2fb750e0ab doc: improve docs for jax.lax trig functions 2025-02-06 11:09:55 -08:00
Michael Hudgins
2e808f2836 Merge pull request #26279 from MichaelHudgins:tsan-resultstore 
PiperOrigin-RevId: 723918760
 2025-02-06 14:55:57 +00:00
Qazalbash
7fc605f783
Merge branch 'main' into scipy-expon 2025-02-05 23:33:51 +05:00
Jake VanderPlas
e4dac395a5 Roll back multinomial change from https://github.com/jax-ml/jax/pull/25688 
This has test breakages on TPU: https://github.com/jax-ml/jax/actions/runs/13159081976/job/36723019653

Reverts 95535df13b422284043623ca3a6d2a5962116fb1

PiperOrigin-RevId: 723536107
 2025-02-05 09:13:56 -08:00
jax authors
6281b86008 Merge pull request #26289 from reikdas:reikdas/docfix 
PiperOrigin-RevId: 722878611
 2025-02-03 18:29:41 -08:00
Pratyush Das
72b6704d88 documentation: fix function name 2025-02-03 19:09:00 -05:00
jax authors
95535df13b Merge pull request #25688 from carlosgmartin:random_multinomial 
PiperOrigin-RevId: 722741835
 2025-02-03 11:52:43 -08:00
jax authors
12c76cdeaa Merge pull request #25474 from jaro-sevcik:compilation-cache-mock-doc 
PiperOrigin-RevId: 722681339
 2025-02-03 09:07:31 -08:00
Qazalbash
46c5b865b7
docs: Update jax.scipy.rst to include exponential distribution functions in documentation 2025-02-01 12:55:49 +05:00
carlosgmartin
32411a430f Add jax.random.multinomial. 2025-01-31 18:45:55 -05:00
Axel Donath
dd7c203b81 Adjust conf.py to remove prompt from sphinx copybutton 2025-01-31 13:51:27 -05:00
Dan Foreman-Mackey
e2eff1f8d5 Revert https://github.com/jax-ml/jax/pull/25982 since callbacks can now use JAX functions. 2025-01-29 11:12:32 -05:00
jax authors
bf22b53cf4 Merge pull request #26154 from jakevdp:pure-callback-doc 
PiperOrigin-RevId: 720763192
 2025-01-28 17:28:02 -08:00
Jake VanderPlas
25aa5a3008 DOC: avoid deprecated argument in external callbacks 2025-01-28 11:34:56 -08:00
Jake VanderPlas
ba2858f834 DOC: add discussion of exceptions in pure_callback 2025-01-28 09:53:47 -08:00
jax authors
6004a501ad Merge pull request #26129 from jakevdp:osx-support 
PiperOrigin-RevId: 720246018
 2025-01-27 11:32:37 -08:00
Jake VanderPlas
7eacce0b83 DOC: update osx x86 entry in installation grid 2025-01-27 11:18:33 -08:00
Dan Foreman-Mackey
782138fb6f Add custom_dce to changelogs and API docs. 2025-01-27 13:03:34 -05:00
jax authors
726abc9c31 Merge pull request #26082 from Saransh-cpp:index-update-syntax-link 
PiperOrigin-RevId: 719356167
 2025-01-24 10:37:07 -08:00
Saransh Chopra
9e89ae7a19 docs: outdated link for index update syntax 2025-01-24 17:36:35 +00:00
Stephan Hoyer
b2afb5bf4f add docstrings for check_vjp and check_jvp 2025-01-23 17:31:13 -08:00
Stephan Hoyer
458f6a6efe Add jax.test_util to public API docs 2025-01-23 16:04:35 -08:00
Jake VanderPlas
423be16ecc DOC: mention adoption of EffVer in JEP 2025-01-22 15:44:11 -08:00
Jakob Roth
4f8699c8a1 Update docs of callbacks 
Callback functions should not call into JAX. This information was
missing in the docs of the callbacks. This commit adds this information
to the docs.

See: #25861, #24255
 2025-01-19 11:33:20 +01:00
Jake VanderPlas
3141453b1b Add performance note at the top of sparse docs 2025-01-17 10:17:16 -08:00
Robert Dyro
aa9cea0a55 Add quotes to pip commands in docs around option install for zsh 
PiperOrigin-RevId: 716358130
 2025-01-16 13:30:43 -08:00
Zac Mustin
2d72e8de84 Jax: Stop returning a list of cost-analyses. 
As it stands, there is only ever one element in this list (see b/384741132) and only the 0th element is ever used so we can simplify.

This is a potentially breaking change for external users, but (as stated in the [documentation](https://jax.readthedocs.io/en/latest/aot.html#debug-information-and-analyses-when-available)) no guarantees are made on this type, which is intended for debugging purposes and not intended to be a reliable public API.

PiperOrigin-RevId: 715837855
 2025-01-15 09:53:59 -08:00
Benjamin Chetioui
57a259f447 [Docs] Remove --xla_gpu_enable_triton_softmax_fusion from docs 
This flag has been a no-op for a while.

PiperOrigin-RevId: 715491248
 2025-01-14 12:50:59 -08:00
Roy Frostig
a60ead6fd1 enable partitionable threefry by default 
PiperOrigin-RevId: 715242560
 2025-01-13 22:46:24 -08:00
Ayaka
9ba1fd2801 [Pallas TPU] Add vector support to pl.debug_print 
PiperOrigin-RevId: 715085454
 2025-01-13 13:22:21 -08:00
jax authors
1fe72ee880 Merge pull request #25771 from dfm:custom-partition-ffi-tutorial 
PiperOrigin-RevId: 714048619
 2025-01-10 07:43:10 -08:00
George Necula
dd0447a7c6 [aot] Add support for as_text(debug_info=True). 
This exposes an easier way to get StableHLO and HLO
with more debugging information (source locations
for StableHLO and metadata for HLO).
 2025-01-10 07:59:56 +02:00
Dan Foreman-Mackey
729418094e Add a discussion of sharding to the FFI tutorial. 2025-01-09 11:24:40 -05:00
Sergei Lebedev
90201ce2b7 Removed leftover mentions of xmap from the code 
PiperOrigin-RevId: 713202387
 2025-01-08 01:39:38 -08:00
Jane Liu
21fb171ef9
Merge branch 'jax-ml:main' into activation-offloading-doc 2025-01-07 21:20:29 -08:00
jax authors
57c2afe7a8 Merge pull request #25441 from Exferro:fixed_advanced_autodiff_doc 
PiperOrigin-RevId: 712929769
 2025-01-07 08:56:05 -08:00