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

21 Commits

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