mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
22,387 Commits 503 Branches 263 Tags
Commit Graph

249 Commits

Author SHA1 Message Date
Roy Frostig
371935cc10 update README and several docs to typed RNG keys 2024-08-11 08:09:47 -07:00
Yash Katariya
7f8a4c84d3 Remove PositionalSharding from distributed array doc 2024-08-07 21:25:24 -07:00
jax authors
8ed94bcfb6 [shard_map docs]: Fix doc typos 
PiperOrigin-RevId: 656265100
 2024-07-25 23:29:55 -07:00
Yash Katariya
0d5dae09ff Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period). 
PiperOrigin-RevId: 655614395
 2024-07-24 10:24:09 -07:00
Dan Foreman-Mackey
705eed3388 Fixing dtype canonicalization in sharp edges tutorial. 
As reported in https://github.com/google/jax/issues/22493, the sharp
edges tutorial doesn't seem to actually enable x64 when it says it does.

Fixes https://github.com/google/jax/issues/22493
 2024-07-22 15:02:02 -04:00
jax authors
fc1e1d4a65 Add freshness metablock to JAX OSS docs. 
PiperOrigin-RevId: 645508135
 2024-06-21 14:50:49 -07:00
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
Jake VanderPlas
10ed827fe9 DOC: replace old tutorials with new content 2024-05-03 12:20:06 -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
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
Jake VanderPlas
a174e9b8b0 DOC: update Colab mention in distributed arrays doc 2024-04-11 09:26:14 -07:00
rajasekharporeddy
d92e735001 Fix Typos in docs 2024-04-11 07:54:13 +05:30
Sai-Suraj-27
29def4eefa Updated all the pre-commit hooks versions. 2024-04-08 00:59:02 +05:30
jax authors
2827564e86 Merge pull request #20204 from mattjj:shmap-tutorial-typos-2 
PiperOrigin-RevId: 615147238
 2024-03-12 13:12:54 -07:00
Matthew Johnson
bf495bf5cb fix typos in shard_map tutorial 2024-03-12 11:58:26 -07:00
Sergei Lebedev
88853c1fe3 Add a deprecation warning to the xmap function and the corresponding tutorial 2024-03-11 21:31:32 +00:00
Roy Frostig
75a53f46e0 update notebooks to new-style typed PRNG keys 2024-03-07 12:40:10 -08:00
jax authors
2cfbd35149 Typo: MT19937 is 2.5 kilobytes, not kilobits. 
PiperOrigin-RevId: 603903613
 2024-02-03 02:03:00 -08:00
Jake VanderPlas
a3abd178f0 DOC: run shard_map docs as part of the doc build 2024-01-22 18:07:01 -08:00
jax authors
f21022b13b Merge pull request #19472 from mattjj:shmap-tutorial-typos 
PiperOrigin-RevId: 600619148
 2024-01-22 17:40:15 -08:00
Matthew Johnson
a61f7807b6 [shard-map] fix typos in docs 2024-01-22 17:29:10 -08:00
8bitmp3
c3010ad026 Add shard_map doc title 2024-01-19 18:23:41 +00:00
Parker Schuh
5cf129d2e9 [shard_map docs]: Fix some typos 
PiperOrigin-RevId: 599656511
 2024-01-18 16:18:14 -08:00
jax authors
853a4fe0dc Merge pull request #19280 from mattjj:shmap-tutorial 
PiperOrigin-RevId: 599644820
 2024-01-18 15:36:28 -08:00
Matthew Johnson
8b219d5f7b [shard-map] add user tutorial 2024-01-18 15:30:13 -08:00
Jake VanderPlas
10eae3f93a CI: update jupytext version 2024-01-09 14:34:21 -08:00
Jake VanderPlas
7d8c358fce Fix wording in TFDS example 2023-11-03 11:31:48 -07:00
Jake VanderPlas
389eb97a7c CI: update pre-commit hooks to latest version 2023-10-30 09:12:24 -07:00
parikshit adhikari
3c283db2fc fixed typo 'primtive' as 'primitive' in How_JAX_primitives_work.ipynb 2023-10-13 21:05:43 +05:45
parikshit adhikari
ba1af0114d fix: typo inside docs/notebooks/How_JAX_primitives_work.md 2023-10-13 08:54:51 +05:45
8bitmp3
7e2e4a08e4 Add jit/shard_map banner to xmap docs 
Add jit/shard_map banner to xmap docs
 2023-10-12 22:06:15 +00:00
Peter Hawkins
d8a0227e86 Simplify the torch data loader collate function using tree_map. 
Fixes https://github.com/google/jax/issues/1004
 2023-10-04 14:59:06 -04:00
Jake VanderPlas
4a5bd9e046 Fix typos across the package 2023-09-22 14:54:31 -07:00
Jake VanderPlas
130a53f2a2 DOC: re-enable execution of thinking_in_jax.ipynb 2023-08-24 09:23:26 -07:00
8bitmp3
03ebdb7454 Add Kaggle TPU VM link, update Distributed Arrays and Auto Parallelization guide 2023-08-21 21:17:02 +00:00
8bitmp3
e0bd8a164d Add Kaggle TPU VM link, update Distributed Arrays and Auto Parallelization guide 2023-08-21 21:09:58 +00:00
Peter Hawkins
2c32660a8f Replace references to DeviceArray with Array. 
A number of stale references are lurking in our documentation.
 2023-08-18 17:46:00 -04:00
Peter Hawkins
47651c6a59 Remove uses of XLA translation rules. 
Remove translation_rule argument to standard_primitive.

PiperOrigin-RevId: 557220350
 2023-08-15 12:53:36 -07:00
Jake VanderPlas
e4701b2451 DOC: callback doc is no longer work-in-progress 2023-08-07 11:16:56 -07:00
Jake VanderPlas
7bb8312f82 CI: update jupytext to v0.14.7 2023-07-24 11:51:45 -07:00
Peter Hawkins
816ba91263 Use lower-case PEP 585 names for types. 
Issue https://github.com/google/jax/issues/16537

PiperOrigin-RevId: 542969282
 2023-06-23 15:12:14 -07:00
Peter Hawkins
34d9f5a9ae Add a CI presubmit that renders the documentation. 2023-06-20 09:29:25 -04:00
Jake VanderPlas
47ae5bddd7 Mark jax.abstract_arrays as deprecated 2023-06-07 23:36:40 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
Matthew Johnson
690071f1de fix custom_jvp docs typo 2023-04-15 14:51:33 -07:00
jax authors
5456152f0b Fix typo: "__W__ - spatial height" -> "__W__ - spatial width" 
PiperOrigin-RevId: 523365127
 2023-04-11 04:37:26 -07:00
Jake VanderPlas
4473ebc9fc Add documentation for callback functions 2023-04-07 07:05:46 -07:00
Jake VanderPlas
c7c9cb652e Sharp bits: refer to ndarray.at in out-of-bound indexing discussion 2023-03-17 13:29:05 -07:00
Peter Hawkins
71f120beed Add "Open in Kaggle" buttons to Jupyter notebooks. 2023-03-01 13:15:42 -05:00