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

282 Commits

Author SHA1 Message Date
Yash Katariya
aa9480a441 Expose get_abstract_mesh via the jax.sharding namespace 
PiperOrigin-RevId: 736972976
 2025-03-14 13:45:32 -07:00
Dougal
e8f43d1cef Explicit sharding docs 2025-03-14 16:33:30 -04:00
jax authors
1a57fdf704 Fix convolution example (kernel should be OIHW, not IOHW). 
PiperOrigin-RevId: 732952185
 2025-03-03 09:32:22 -08:00
Jake VanderPlas
f5ca46f5ec Sharp bits: add note on subnormal flush-to-zero 2025-02-24 18:52:38 -08:00
Saransh Chopra
9e89ae7a19 docs: outdated link for index update syntax 2025-01-24 17:36:35 +00:00
Jane Liu
77c6947a59 fix the doc error: module 'scipy.misc' has no attribute 'face' 2025-01-06 12:00:02 -08:00
Matthew Johnson
c6576827c6 remove whitespace from distributed arrays doc 2024-12-18 23:55:20 +00:00
jax authors
cd7109e6b5 Minor typo fixin doc. 
PiperOrigin-RevId: 706774033
 2024-12-16 11:22:45 -08:00
nireekshak
f43fa9fc78 Fix some typos 2024-12-03 05:07:53 +00:00
nireekshak
1458d3dd56 Fix some typos 2024-11-19 15:04:55 +00:00
Emily Fertig
225a2a5f8b Consolidate material on PRNGs and add a short summary to Key Concepts. 2024-11-15 14:44:57 -08:00
Emily Fertig
e6f6a8af8d Move Control Flow text from Sharp Bits into its own tutorial. 2024-11-14 11:07:52 -08:00
Yash Katariya
ca2d1584f8 Remove mesh_utils.create_device_mesh from docs 
PiperOrigin-RevId: 687695419
 2024-10-19 15:48:42 -07:00
Jake VanderPlas
e1f280c843 CI: enable additional ruff formatting checks 2024-10-16 16:09:54 -07:00
Ayaka
e79d77aa47 [Pallas] [Docs] Replace full urls with label-based cross references 
This PR uses the same method to add cross references as the previous PR https://github.com/jax-ml/jax/pull/23889.

---

The content below is for future references.

#### Useful commands

Build documentation:

```sh
sphinx-build -b html -D nb_execution_mode=off docs docs/build/html -j auto
```

Create a label in *.md:

```md
(pallas_block_specs_by_example)=
```

Create a label in *.rst:

```rst
.. _pallas_tpu_noteworthy_properties:
```

Reference a label in *.md:

```md
{ref}`pallas_block_specs_by_example`
```

Sync changes from *.md to *.ipynb:

```sh
jupytext --sync docs/pallas/tpu/distributed.md
```

PiperOrigin-RevId: 682034607
 2024-10-03 14:35:51 -07:00
8bitmp3
0cf040c9a1 Add/update JAX Advanced Tutorials docs, ToC structure 2024-09-20 23:06:54 +00:00
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax 
PiperOrigin-RevId: 676843138
 2024-09-20 07:52:33 -07:00
Roy Frostig
7b29791200 docs: shorten/clarify some page titles 
Shorten the titles for derivative rules and manual parallelism, and go
for canonical wording in the parallel programming intro (typically we
"shard" data, and "partition" computation, as part of parallel
programming).
 2024-08-28 11:22:06 -07:00
Roy Frostig
763d600508 docs: remove inline authors/dates on misc doc pages 
Git history covers this better and automatically.
 2024-08-28 11:04:14 -07:00
Roy Frostig
672a013b3a docs: prefer "summary" to "tl;dr" 
This was common and we typically just mean "summary."
 2024-08-28 11:04:14 -07:00
Jake VanderPlas
09fd345de9 pre-commit: update hooks & pin using hashes 2024-08-27 15:23:13 -07:00
Jake VanderPlas
68be5b5085 CI: update ruff to v0.6.1 2024-08-27 14:54:11 -07:00
carlosgmartin
f812d0f28b Clarify meaning of tall and wide Jacobian matrices in autodiff docs. 2024-08-26 16:00:36 -07:00
jax authors
3e764f617a [shard_map docs]: Fix doc typos 
PiperOrigin-RevId: 664960613
 2024-08-19 13:45:46 -07:00
Roy Frostig
d17edb4c4d docs: fix shard_map guide headings 
These were off by one level, causing section titles to be listed in
the guide index.
 2024-08-13 18:29:42 -07:00
jax authors
d2b85a48af Merge pull request #23036 from froystig:docs2 
PiperOrigin-RevId: 662635310
 2024-08-13 13:40:19 -07:00
Roy Frostig
3c223cd253 docs: tidy up titles and headings 
This shortens some titles and makes them more consistent. It also
removes "JAX" from several titles ("in JAX", "for JAX", "JAX's",
etc.). Since these are JAX docs, that ought to be clear from context.
 2024-08-13 11:53:57 -07:00
jax authors
23effba503 Merge pull request #23027 from froystig:docs 
PiperOrigin-RevId: 662392585
 2024-08-12 23:59:39 -07:00
Roy Frostig
09e73118bf docs: more sentence case 2024-08-12 20:07:49 -07:00
jax authors
833560deb1 Merge pull request #23023 from froystig:docs3 
PiperOrigin-RevId: 662318149
 2024-08-12 19:03:12 -07:00
Roy Frostig
b8f8b7b07f docs: sentence case page titles, section headings, some content 2024-08-12 18:12:17 -07:00
jax authors
312ebdb8ba Merge pull request #22989 from froystig:keys 
PiperOrigin-RevId: 661890876
 2024-08-11 16:08:25 -07:00
Roy Frostig
371935cc10 update README and several docs to typed RNG keys 2024-08-11 08:09:47 -07:00
Roy Frostig
ded5b5366b indent consistently in auto-parallelization and shard_map tutorials 2024-08-11 08:04:14 -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