mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 15:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
21,905 Commits 508 Branches 263 Tags
Commit Graph

1804 Commits

Author SHA1 Message Date
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
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