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

29 Commits

Author SHA1 Message Date
Justin Fu
dcd186f552 [Pallas] Add pallas distributed computation tutorial 2024-08-08 17:34:35 -07:00
Justin Fu
f8919a32e0 Fix minor typo in Pallas docs. 
PiperOrigin-RevId: 625117045
 2024-04-15 16:20:42 -07:00
jax authors
51352fa05c fix matrix dimension and block shape. 
PiperOrigin-RevId: 624988654
 2024-04-15 09:39:31 -07:00
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
Sai-Suraj-27
29def4eefa Updated all the pre-commit hooks versions. 2024-04-08 00:59:02 +05:30
Sergei Lebedev
ea8e393c0e Fixed a few typos in the matmul example in "Pallas Design" 2024-04-03 10:46:05 +01:00
Sharad Vikram
87aee90e67 Fix typo in Pallas design 
PiperOrigin-RevId: 621275025
 2024-04-02 13:20:46 -07:00
rajasekharporeddy
61c64c10f8 Fixed Several Typos 
Fixed Typos in JEP doc files

Revert "Fixed Typos in JEP doc files"

This reverts commit c2a16950e0fc1b32971168501d183991e2394b5d.

revert two changes

reverted one change in advanced-autodiff

revert one change in parallelism

sync notebooks
 2024-03-12 00:37:46 +05:30
Roy Frostig
fe3e798f82 update pallas quickstart to new-style typed PRNG keys 2024-03-07 12:40:10 -08:00
Sharad Vikram
30973a9474 [Pallas] Pass in compiler params via explicit compiler_params argument instead of passing via **kwargs 
This is a change that makes the API a bit more intuitive and avoids footguns like accidentally passing in `in_spec` instead of `in_specs` because previously kwargs that weren't used by any downstream lowering would be ignored and users would get weird errors as a result.

This change doesn't deprecate the old way of passing in compiler params but it will be deprecated soon after this.

PiperOrigin-RevId: 613239439
 2024-03-06 09:16:22 -08:00
Sharad Vikram
2f8d5cebff [Pallas TPU] Add Pipelining and BlockSpecs documentation for Pallas TPU 2024-02-28 19:17:58 -08:00
jax authors
576b54661d Minor typos in pallas documentation 
PiperOrigin-RevId: 606230620
 2024-02-12 06:19:35 -08:00
Peter Hawkins
e558feaa5e Deprecate support for the mhlo dialect. 
JAX has not used mhlo for some time, in favor of stablehlo. Deprecate support for this dialect in JAX's API and remove testing.

PiperOrigin-RevId: 598550225
 2024-01-15 02:13:40 -08:00
Matthew Johnson
efe78c53ec improve block matrix alignment in pallas docs 2024-01-10 22:01:16 -08:00
jax authors
88169cf9e5 Merge pull request #19275 from j2kun:main 
PiperOrigin-RevId: 597148779
 2024-01-09 23:06:02 -08:00
Jeremy Kun
2e6e5da49b docs/pallas: remove list from out_specs 2024-01-09 15:47:18 -08:00
Jeremy Kun
4ecbed1322 docs/pallas: sync jupyter notebook 2024-01-09 15:00:42 -08:00
Jake VanderPlas
10eae3f93a CI: update jupytext version 2024-01-09 14:34:21 -08:00
Jeremy Kun
87f914e02f docs/pallas: fix mismatched pytree specs 
Running the example as-is gives

```
ValueError: Pytree specs for `out_shape` and `out_specs` must match: PyTreeDef(*) vs. PyTreeDef((*,))
```

Giving a list argument to `out_shape` seems to fix the issue.
 2024-01-09 13:56:28 -08:00
Alexey Radul
1f184541e4 Ah, jupytext converts from markdown to a notebook, not vice versa. 2023-11-09 11:11:21 -05:00
Alexey Radul
785d9fa14c Lightly proofread Pallas documentation. 2023-11-09 10:57:47 -05:00
Jake VanderPlas
389eb97a7c CI: update pre-commit hooks to latest version 2023-10-30 09:12:24 -07:00
Jake VanderPlas
4a5bd9e046 Fix typos across the package 2023-09-22 14:54:31 -07:00
Sharad Vikram
00bc7cf6f7 [Pallas] Fix out_specs in Pallas quickstart 2023-09-21 17:58:57 -07:00
jax authors
22285e69fb Merge pull request #16971 from apaszke:pallas-tpu-docs 
PiperOrigin-RevId: 554587987
 2023-08-07 14:10:07 -07:00
Sharad Vikram
28af1861ee [Pallas] Fix rendering of math in quickstart 2023-08-05 00:20:00 -07:00
Adam Paszke
98191dab75 Add a guide for writing Pallas TPU kernels 2023-08-04 14:27:58 +00:00
Sharad Vikram
96e2d93f53 [Pallas] Add Pallas design doc 2023-08-03 16:02:07 -07:00
Sharad Vikram
7f1ef32ba3 Add initial documentation for Pallas 2023-08-03 12:30:19 -07:00