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

1818 Commits

Author SHA1 Message Date
George Necula
0223e830c1 [docs] Enable doctest for .md documentation files 
The current invocation of doctest from GH actions picked up only .rst files.
We enable .md files also, and we make a few changes to ensure that the
doctest passes on the existing files.

The changes fall into several categories:
  * add a newline before the end of the code block, for doctest to
    pick up the expected output properly
  * update the expected values to match the current behavior
  * disable some doctests that raise expected exceptions, whenever
    I could not get doctest to match the exception details.
    Sometimes +IGNORE_EXCEPTION_DETAIL was enough, and other times
    I had to use +SKIP.
 2024-06-03 11:27:34 +03:00
Marco Selvi
7a8bcf6ee5
Fix parenthesis in "Gradients contain NaN where using where" 2024-05-31 12:03:39 +01:00
jax authors
ef0b5d7385 Merge pull request #21442 from vfdev-5:added-trace-alias-to-linalg 
PiperOrigin-RevId: 638477013
 2024-05-29 18:27:20 -07:00
vfdev-5
d2185d3636 Added trace alias to jnp.linalg 
Related to #21088
 2024-05-29 22:28:44 +00:00
Frederic Bastien
6466154386 Add a link to another page with XLA flags. 2024-05-29 14:29:24 -04:00
Sam Ritchie
822bc3647d
Update quickstart to suggested jax[cuda12] installation 
This PR updates the quickstart docs to match the README and Installation Guide suggestion of using

```
pip install jax[cuda12]
```

for GPU.
 2024-05-28 13:31:14 -04:00
Yash Katariya
08b1cef00a Add jax.make_array_from_process_local_data to the docs. Preview: https://jax--21444.org.readthedocs.build/en/21444/_autosummary/jax.make_array_from_process_local_data.html#jax.make_array_from_process_local_data 
PiperOrigin-RevId: 637660828
 2024-05-27 09:41:43 -07:00
Jake VanderPlas
1327143d46 Better documentation for jax.tree_util 2024-05-20 19:56:47 -07:00
jax authors
bb616eff8a Merge pull request #21231 from nouiz:doc_experimental_serialize_executable 
PiperOrigin-RevId: 635428472
 2024-05-20 05:39:21 -07:00
Dan Foreman-Mackey
be0695474a Add RegularGridInterpolator to generated API docs 
In responding to gh21279, I noticed that `RegularGridInterpolator` isn't
currently listed in the API docs. I know that `scipy.interpolate` is out
of scope (https://jax.readthedocs.io/en/latest/jep/18137-numpy-scipy-scope.html#scipy-interpolate),
but since we do currently provide this wrapper, it seems like it makes
sense to include it in the docs!
 2024-05-17 11:24:35 -04:00
jax authors
1829a66739 Merge pull request #21268 from jakevdp:register-dataclass 
PiperOrigin-RevId: 634624518
 2024-05-16 21:27:30 -07:00
jax authors
8493fc975d Merge pull request #21181 from 8bitmp3:add-distributed-data-loading-doc 
PiperOrigin-RevId: 634595201
 2024-05-16 18:37:14 -07:00
Jake VanderPlas
defb53f446 Document jax.tree_util.register_dataclass 2024-05-16 18:34:54 -07:00
8bitmp3
15a3debee8 Add Distributed data loading doc 2024-05-16 23:12:13 +00:00
Sergei Lebedev
01194bd2fb Clarified the type of the inputs to callback APIs 
The callback APIs were migrated to use jax.Arrays for both inputs and outputs
in JAX 0.4.27.

PiperOrigin-RevId: 634473890
 2024-05-16 11:29:09 -07:00
Vadym Matsishevskyi
517e299a9d Use hermetic Python in JAX, see "Managing hermetic Python" in developer.md for details 
PiperOrigin-RevId: 634146391
 2024-05-15 18:20:56 -07:00
Frederic Bastien
936ec98380 Add doc for jax.experimental.serialize_executable module 2024-05-14 19:11:37 -04:00
jax authors
9af0a317a5 Merge pull request #21222 from tempoxylophone:patch-1 
PiperOrigin-RevId: 633734390
 2024-05-14 16:05:17 -07:00
Jake VanderPlas
bb5787da09 Finalize deprecations of several APIs 
PiperOrigin-RevId: 633634215
 2024-05-14 10:40:40 -07:00
tempoxylophone
4bcad15111
Fix type in advanced-autodiff.md 
removed `doc.g` text from the first section.
 2024-05-14 11:52:41 -04:00
Jake VanderPlas
72e9eb9367 shard_map: add API docs 2024-05-13 13:04:15 -07:00
jax authors
979d9ca3e5 Merge pull request #21168 from 8bitmp3:upgrade-sharded--doc 
PiperOrigin-RevId: 632648408
 2024-05-10 17:44:15 -07:00
8bitmp3
9ea3fcbab3 Upgrade JAX Parallelism Sharded Computation 101 doc 2024-05-10 21:24:16 +00:00
djm3622
aa48a55173
Update line127 error, debugging.md 
There is an error with "y, z = jnp.sin(x), jnp.cos(x)" where jnp.cos(x) was nested within jnp.sin(x) ==> jnp.sin(x, jnp.cos(x)). This caused an error to be thrown. This change fixes that.
 2024-05-09 22:44:27 -04:00
rajasekharporeddy
0a6827f4e3 Add sparse_sigmoid to jax.nn docs 2024-05-08 12:23:01 +05:30
Jieying Luo
5ba56bb075 Recommend the plugin in the CUDA installation instructions. 
PiperOrigin-RevId: 631555876
 2024-05-07 14:47:39 -07:00
Jake VanderPlas
09810be0cd Implement jnp.linalg.multi_dot using opt_einsum 2024-05-07 13:40:25 -07:00
Samuel
fdd8b61094
Minor doc fix to jaxpr.rst 
Update the written description of the example `jaxpr` to match the given snippet of code.
 2024-05-07 17:36:19 +01:00
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