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

1690 Commits

Author SHA1 Message Date
Jake VanderPlas
10ed827fe9 DOC: replace old tutorials with new content 2024-05-03 12:20:06 -07:00
Stephan Hoyer
c77370c492 Recursively pull out __wrapped__ in linkcode_resolve() 
This should actually fix the source lookup for `jax.nn.relu`, which
uses both `custom_jvp` and `jit` decorators.
 2024-05-03 09:10:00 -07:00
Stephan Hoyer
f6ee61273a Make linkcode_resolve() a bit more robust 
This handles three cases that came up when adopting this snippet (for
finding source code corresponding to API docs) [for neuralgcm](https://github.com/google-research/neuralgcm/pull/58):

1. documenting class method or attributes
2. documenting properties
3. documenting `jit` decorated methods

I'm not sure if case (1) or (2) comes up in the JAX docs, but case (3)
definitely does -- `jit` decorated functions like `jax.nn.relu`
[currently do not](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.relu.html)
have source code link.
 2024-05-02 17:57:02 -07:00
Jake VanderPlas
18e4cfa911 DOC: Improve remaining jax.scipy docstrings 2024-05-02 10:29:00 -07:00
jax authors
f4cf69f118 Merge pull request #20936 from jakevdp:py-support 
PiperOrigin-RevId: 630054494
 2024-05-02 07:17:45 -07:00
Jake VanderPlas
18703d9385 DOC: Improve docstrings for jax.scipy.signal 2024-05-01 14:26:00 -07:00
Jake VanderPlas
1936317b8b jax.scipy.stats: manually document functions to avoid scipy import 2024-05-01 09:26:26 -07:00
jax authors
17d0eb30fa Merge pull request #20961 from jakevdp:linkcode 
PiperOrigin-RevId: 629568292
 2024-04-30 16:56:14 -07:00
Jake VanderPlas
55e3be65b7 Remove dependency on sphinx_autodoc_typehints 2024-04-29 14:19:46 -07:00
Jake VanderPlas
74f1d8897c DOC: add manual documentation to jax.scipy.special functions. 
This lets us give more implementation-specific information, and
lets us avoid a needless dependency on scipy.
 2024-04-29 10:58:07 -07:00
Jake VanderPlas
8915ed8348 DOC: link docstrings to github source 2024-04-26 15:56:52 -07:00
Jake VanderPlas
51fa985050 API deprecation doc: explicitly list some dates 2024-04-25 09:21:16 -07:00
Dan Zheng
e9e3c80258 [jaxpr.rst] Remove extraneous 'let' in Jaxpr grammar. 
Remove `let` from `Eqn` rule. `let` appears only once in the `Jaxpr` rule.

Consistently capitalize grammar nonterminals: `jaxpr` → `Jaxpr`.
 2024-04-25 09:53:16 -04:00
Jake VanderPlas
a8af2b788a DOC: respond to mattjj comments 2024-04-23 13:15:20 -07:00
Matthew Johnson
94e3a6efae quickstart tweaks (from https://github.com/google/jax/pull/20819) 2024-04-19 12:23:26 -07:00
Jake VanderPlas
568db105ea Add jax.scipy.special.gammasgn 2024-04-18 16:14:55 -07:00
Jake VanderPlas
48e8457b81 DOC: one last readthrough of the new 101 tutorials 2024-04-17 16:08:38 -07:00
jax authors
1c8534e034 Merge pull request #20802 from jakevdp:dup-thinking-in-jax 
PiperOrigin-RevId: 625738710
 2024-04-17 10:55:11 -07:00
Jake VanderPlas
edfdc36b4c DOC: remove copy of thinking-in-jax from new tutorial flow 2024-04-17 10:27:15 -07:00
Jake VanderPlas
29b15e2e2d DOC: pin sphinx to >=7.3.2 2024-04-17 09:51:27 -07:00
jax authors
1d615a3bbd Merge pull request #20792 from jakevdp:fix-sharding-doc 
PiperOrigin-RevId: 625515766
 2024-04-16 18:32:02 -07:00
jax authors
fc654a5772 Merge pull request #20732 from jakevdp:doc-stateful 
PiperOrigin-RevId: 625515752
 2024-04-16 18:27:26 -07:00
jax authors
07063fb11f Merge pull request #20794 from jakevdp:pin-sphinx 
PiperOrigin-RevId: 625515741
 2024-04-16 18:27:05 -07:00
jax authors
e554416d2c Merge pull request #20784 from jakevdp:doc-installation 
PiperOrigin-RevId: 625515725
 2024-04-16 18:22:29 -07:00
Jake VanderPlas
3c18a021c3 DOC: add stateful computation doc 2024-04-16 17:54:55 -07:00
Jake VanderPlas
40ad5b283f Pin sphinx version to avoid error in 7.3.0 2024-04-16 16:51:06 -07:00
Jake VanderPlas
2552c29028 DOC: fix Mesh construction in sharding doc 2024-04-16 16:26:18 -07:00
Meekail Zain
ceeb975735 Add new cumulative_sum function to numpy and array_api 2024-04-16 19:57:55 +00:00
Jake VanderPlas
a4ccb354f8 DOC: update installation guide 2024-04-16 12:46:36 -07:00
Yue Sheng
1a650cdc00 Update JAX official doc: point out that the device numbers are not in numerical order because of the underlying torus hardware topology. 
PiperOrigin-RevId: 625401373
 2024-04-16 11:37:19 -07:00
Jake VanderPlas
8e34da70f8 DOC: add introduction to sharded computation 2024-04-16 10:39:18 -07:00
Justin Fu
f8919a32e0 Fix minor typo in Pallas docs. 
PiperOrigin-RevId: 625117045
 2024-04-15 16:20:42 -07:00
Meekail Zain
6bdc83c680 Add new unstack function to numpy/array_api namespaces 2024-04-15 21:03:26 +00:00
jax authors
51352fa05c fix matrix dimension and block shape. 
PiperOrigin-RevId: 624988654
 2024-04-15 09:39:31 -07:00
jax authors
198ce52bf7 Merge pull request #20736 from jakevdp:doc-strict-promotion 
PiperOrigin-RevId: 624319591
 2024-04-12 16:08:45 -07:00
Jake VanderPlas
2be17dc778 DOC: document strict dtype promotion mode 2024-04-12 14:05:05 -07:00
jax authors
7fb27a4cce Merge pull request #19604 from Micky774:faq_cuda 
PiperOrigin-RevId: 624234866
 2024-04-12 11:09:12 -07:00
Jake VanderPlas
f090074d86 Avoid 'from jax import config' imports 
In some environments this appears to import the config module rather than
the config object.
 2024-04-11 13:23:27 -07:00
jax authors
109f5f93aa Merge pull request #20713 from jakevdp:parallelism-doc 
PiperOrigin-RevId: 623877877
 2024-04-11 11:09:08 -07:00
jax authors
11fc7a38e0 Merge pull request #20687 from rajasekharporeddy:test_branch4 
PiperOrigin-RevId: 623877846
 2024-04-11 11:08:53 -07:00
Jake VanderPlas
a174e9b8b0 DOC: update Colab mention in distributed arrays doc 2024-04-11 09:26:14 -07:00
Jesse Rosenstock
69be3b207c
DOC: Fix :ref: format 
Change ref format from {ref}`working-with-pytrees` to
:ref:`working-with-pytrees`.  The old version just renders literally, 
not as a link.  It appears to be a markdown-style reference.  I'm
not familiar with reST, but :ref: is how the other references are
formatted.

Introduced in d1235fa (DOC: add key concepts doc).
 2024-04-11 12:54:37 +02:00
rajasekharporeddy
d92e735001 Fix Typos in docs 2024-04-11 07:54:13 +05:30
Sergei Lebedev
a205c9120a pallas_call now has only one way to pass compiler_params= 
Previously, it was possible to do

    pallas_call(..., foo=42)

and also

    pallas_call(..., compiler_params=dict(foo=42))

PiperOrigin-RevId: 623277572
 2024-04-09 14:23:20 -07:00
jax authors
afb775c168 Jax persistent compilation cache user guide. 
This user guide covers using the cache on local filesystems
and Google Cloud.

PiperOrigin-RevId: 623236335
 2024-04-09 11:48:33 -07:00
jax authors
f380c3c365 Merge pull request #20179 from keshavb96:custom_gpu_ops_partitioning 
PiperOrigin-RevId: 622909424
 2024-04-08 11:42:49 -07:00
Frederic Bastien
83ffcc9c7d Current status + build script fixes 
Add print

First version with custom_partitioning. The communication during the gradient aren't optimal.

Fix the gradient sharding

small update

Fix the strange replicated computation.

Make it work with the new JAX version.

Add the structure for custom_p domentation.

Small clean up

First version of the doc

Add comment and typing annotation

tab->space

Simplify code and add docstring

Use the simpler JAX API since 0.4.16 (August 2023).

Custom partitioning using custom_partitioning

updated docs; dump custom_partitioning HLO

doc update

more documentation updates; include links to code instead of inlined code

fix typos

fix more typos

fix type annotations in source and update docs

minor fixes

import fix

lint fix

added apache license header
 2024-04-08 04:48:28 -07:00
Sai-Suraj-27
29def4eefa Updated all the pre-commit hooks versions. 2024-04-08 00:59:02 +05:30
Meekail Zain
6b97a319e6 Updated FAQ to include section on CUDA library initialization failures 2024-04-04 23:08:10 +00:00
jax authors
29a2762b64 Merge pull request #20558 from carlosgmartin:mish 
PiperOrigin-RevId: 621708823
 2024-04-03 19:50:58 -07:00