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

25 Commits

Author SHA1 Message Date
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
jax authors
fc1e1d4a65 Add freshness metablock to JAX OSS docs. 
PiperOrigin-RevId: 645508135
 2024-06-21 14:50:49 -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
rajasekharporeddy
75da54e9f4 Fix doc typos 2024-03-13 14:48:31 +05:30
Jake VanderPlas
2f621119ea DOC: mention static values in debug.print docs 2023-08-28 10:59:40 -07:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07:00
Skye Wanderman-Milne
473d1c3685 Turn on PJRT C API by default. 
I forgot that the default setting is actually in jaxlib:
fbe9a80fdb/xla/python/xla_client.py (L135)

To be able to make this change as a jax-only release, I manually set
the env var on Cloud TPU if it isn't already set.
 2023-03-28 15:28:13 -07:00
Matthew Johnson
1f67351f56 [shard_map] make debug_print work with shard_map, eager and jit 2023-03-08 20:38:03 -08:00
Yash Katariya
52a7701dda Replace usage of {in|out}_axis_resources with {in|out}_shardings 
PiperOrigin-RevId: 513040164
 2023-02-28 14:29:09 -08:00
Sharad Vikram
6f1714e57a Add some info in the docs about using jax.debug.print with f-strings 2023-02-15 15:16:37 -08:00
Peter Hawkins
428189f8fb Replace uses of deprecated JAX sharding APIs with their new names in jax.sharding. 
This change updates:
* {jax.experimental.maps.Mesh, jax.interpreters.pxla.Mesh} to jax.sharding.Mesh
* {jax.experimental.PartitionSpec, jax.experimental.pjit.PartitionSpec, jax.interpreters.pxla.PartitionSpec, jax.pxla.PartitionSpec} to jax.sharding.PartitionSpec
* jax.experimental.maps.NamedSharding to jax.sharding.NamedSharding.

PiperOrigin-RevId: 506994892
 2023-02-03 14:28:45 -08:00
Rafi Witten
278ff25ae1 Update docs that jax.debug is unsupported on Cloud TPUs 2023-02-01 00:12:51 +00:00
8bitmp3
d6cc2bdb22 Make JAX Debugging and Profiling guides more visible, move Profiling to User Guides from Notes 2023-01-19 23:05:24 +00:00
lenamartens
8a4b9d6aad Fix typo in checkify guide. 2023-01-05 17:04:53 +00:00
lenamartens
3db909ee8d Checkify: Add documentation for adding run-time values to error message. 2022-12-12 15:56:40 +00:00
jax authors
cdd16e3bbd Fix typo in jax.debug.print() documentation. 
PiperOrigin-RevId: 471096544
 2022-08-30 14:51:12 -07:00
Sharad Vikram
27b313b287 Fix typo in debugging docs 2022-08-29 10:32:08 -07:00
Sharad Vikram
8b54298052 Fix typo in debugging docs 2022-08-12 11:59:00 -07:00
Sharad Vikram
11b206a18a Enable debugging primitives in pjit on CPU/GPU 
PiperOrigin-RevId: 464208326
 2022-07-29 20:10:27 -07:00
jax authors
6cffa720e7 Merge pull request #11670 from sharadmv:debugging-docs 
PiperOrigin-RevId: 464145832
 2022-07-29 13:26:56 -07:00
Sharad Vikram
decdca60c8 Change jaxdb->jdb and add option to force a backend 2022-07-29 12:51:27 -07:00
jax authors
cc19c94d36 Merge pull request #11665 from LenaMartens:docs 
PiperOrigin-RevId: 464113785
 2022-07-29 10:56:58 -07:00
lenamartens
0256301617 tweak checkify docs. 2022-07-29 18:15:36 +01:00
Sharad Vikram
fb0cf668b8 Update debugging docs to mention pjit 2022-07-28 22:00:27 -07:00
Sharad Vikram
4386a0f909 Add debugging tools under jax.debug and documentation 
Co-authored-by: Matthew Johnson <mattjj@google.com>
Co-authored-by: Lena Martens <lenamartens@google.com>
 2022-07-28 20:07:26 -07:00