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

16 Commits

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