Sharad Vikram
|
2d8b228706
|
Add function to visualize Shardings
|
2022-09-19 13:27:08 -07:00 |
|
Sharad Vikram
|
7cd81ca1a8
|
Allow debug prints in staged out custom derivative functions
PiperOrigin-RevId: 467344265
|
2022-08-12 19:49:09 -07:00 |
|
Sharad Vikram
|
2d72dc899b
|
Enable receives in TPU callbacks and add tests
PiperOrigin-RevId: 466103306
|
2022-08-08 11:42:16 -07:00 |
|
Sharad Vikram
|
7c49f9a603
|
Error in debug print if unused arguments are provided
|
2022-08-05 17:24:04 -07:00 |
|
Sharad Vikram
|
89150eef1d
|
Enable debug callbacks in checkpoint
PiperOrigin-RevId: 465601044
|
2022-08-05 10:59:32 -07:00 |
|
jax authors
|
5ec9226845
|
Merge pull request #11717 from sharadmv:move-sharding
PiperOrigin-RevId: 465369371
|
2022-08-04 11:55:57 -07:00 |
|
Sharad Vikram
|
5f618e706d
|
Throw error earlier for misformatted string in jax.debug.print
|
2022-08-03 19:03:57 -07:00 |
|
Sharad Vikram
|
375ef0bc63
|
Making sharding an arg to MHLO callback lowering
|
2022-08-03 11:04:14 -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 |
|
Sharad Vikram
|
fc1fa134c8
|
Adjust debug_callback JVP rule to only call on primals
|
2022-07-22 15:47:23 -07:00 |
|
Peter Hawkins
|
0b4b0ba072
|
Update minimum jaxlib version to 0.3.14.
|
2022-07-08 00:36:02 +00:00 |
|
Sharad Vikram
|
6274b9ed39
|
Enable Python callbacks on TFRT TPU backend
PiperOrigin-RevId: 459415455
|
2022-07-06 20:52:50 -07:00 |
|
Sharad Vikram
|
b80d7195f6
|
Enable python callback on GPU
PiperOrigin-RevId: 452354891
|
2022-06-01 12:15:07 -07:00 |
|
Sharad Vikram
|
4296cb1c2b
|
Enable AD rules for debug_print
|
2022-05-05 10:11:58 -07:00 |
|
Sharad Vikram
|
c1a8d7fb57
|
Enable batching rule for debug_print
Co-authored-by: Matthew Johnson <mattjj@google.com>
|
2022-05-04 23:32:36 -07:00 |
|
Sharad Vikram
|
dc42d7bb8e
|
Enable print in while loops/scan
|
2022-05-04 18:55:47 -07:00 |
|
Sharad Vikram
|
78a4e30651
|
Don't leak the keepalive in debug_callback lowering
|
2022-05-03 20:53:30 -07:00 |
|
Sharad Vikram
|
7a641de7a7
|
Add initial debug print implementation
Co-authored-by: Matthew Johnson <mattjj@google.com>
|
2022-05-03 20:16:05 -07:00 |
|