mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
13,057 Commits 506 Branches 263 Tags
Commit Graph

35 Commits

Author SHA1 Message Date
lenamartens
018e700ead Checkify: support batched while. 2022-09-20 17:59:46 +01:00
Lena Martens
b3393e3b60 Checkify: add jvp rule for assert_p. 
PiperOrigin-RevId: 472961963
 2022-09-08 05:26:02 -07:00
Yash Katariya
acdae7c237 Add weak type support to Array. Also make all api_test.py tests pass with Array. I have disabled the float0 test for now until I investigate. 
PiperOrigin-RevId: 468264910
 2022-08-17 12:25:49 -07:00
Yash Katariya
4fc3518e5f Make checkify tests pass with Array and add methods on Array that are present on DA. 
PiperOrigin-RevId: 468058909
 2022-08-16 16:52:06 -07:00
lenamartens
1ace5d351b Checkify: support checkify-of-pjit. 2022-07-29 19:25:22 +01:00
Lena Martens
740fe6926a Checkify: add (checkify-of-)vmap-of-check. 2022-06-27 10:34:26 +01:00
Lena Martens
9167f7248a Checkify: support discharging checks from control-flow through effects. 
Currently supports scan and while-loop.
 2022-06-20 18:28:03 +01:00
Jake VanderPlas
305f5e0491 [x64] make checkify compatible with strict dtype promotion 2022-06-17 12:48:20 -07:00
Jake VanderPlas
d2f80ef117 [x64] deprecate unsafe type casting in scatter-update operations 2022-06-09 15:21:49 -07:00
Jeppe Klitgaard
17de89b16a feat: refactor code using pyupgrade 
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```

a
 2022-05-17 22:14:05 +01:00
Lena Martens
e4836f5663 Checkify: support checks on data-independent values. 
You can now check values which do not depend on checkified args.

Co-authored-by: Matthew Johnson <mattjj@google.com>
 2022-04-19 15:06:00 +01:00
Lena Martens
15d2ccaeba Checkify: add axis and axis size to OOB error message. 2022-03-31 15:16:35 +01:00
Lena Martens
d72687c990 Checkify: add way to embed runtime info in error messages. 2022-03-28 11:55:34 +01:00
Lena Martens
76e1021cad Checkify: Fix empty enabled_errors case wrt user checks. 
What looked like a quick win (short-cut on empty enabled_error, don't trace)
was actually a quick bug (user_checks throw error in eager).
 2022-03-10 13:33:29 +00:00
lenamartens
45d3ddda31 Fix tests and handle cond consts. 2022-02-23 16:11:09 +00:00
Matthew Johnson
4b1d0a466b fixing scan and other control flow 
Co-authored-by: Lena Martens <lenamartens@google.com>
 2022-02-23 15:39:22 +00:00
Lena Martens
2eeb683df0 Checkify: address initial feedback. 
- add a way to run checkify with no errors enabled
  - clarify "can't be staged" error message
  - export init_error: a way for users to set a default Error value
 2022-02-18 19:23:28 +00:00
Lena Martens
758c721605 Checkify: fix nd-error case when array only has 1 element. 2022-02-16 19:15:48 +00:00
Lena Martens
b15c7f609a Checkify: fix check_error of nd-error. 
PiperOrigin-RevId: 428857813
 2022-02-15 13:12:53 -08:00
Lena Martens
a4cacf5729 Checkify: handle named_call in process_call. 
named_call does not specify donated_invars, this change handles this missing
param case.

For future reference: we might want to add a call_param_updater registry to define
how call params need to get updated wrt checkify, like eg. partial_eval/ad does.
 2022-02-14 15:24:55 +00:00
Lena Martens
2ba8aec274 Checkify: Fix docstring formatting and polish enabled_errors sets. 2022-02-10 16:54:43 +00:00
Matthew Johnson
d9270b242d [checkify] add custom_vjp support 2022-02-09 12:31:16 -08:00
Matthew Johnson
4ce749e681 [checkify] handle custom_jvp 2022-02-09 12:12:58 -08:00
Lena Martens
0042edb5f4 Checkify: rename some symbols and add some docstrings. 2022-02-08 17:40:04 +00:00
Matthew Johnson
9488c5ae72 checkify: fix and test post_process_call/map 2022-01-19 14:44:30 -08:00
Lena Martens
8ea85769ea Checkify: add way to disable categories of errors. 
By default only user_asserts are lifted into the checked function.
 2022-01-18 17:59:50 +00:00
Lena Martens
7b5b9cefbd Add scatter OOB error. 2022-01-07 17:22:34 +00:00
Lena Martens
bbd127f8fa Add division by zero check. 2022-01-04 19:17:08 +00:00
Lena Martens
98a5461132 Make sure while_loop cond generates an error even if it returns False. 2021-12-16 21:50:45 +00:00
Lena Martens
0dc5a33a88 Add checkify rule for while_loop. 2021-12-08 19:34:20 +00:00
Lena Martens
bbf1a9ba97 Add checkify rule for scan. 2021-12-07 12:36:44 +00:00
Matthew Johnson
c1f71d17c0 generalize assert primitive, allow recharging 2021-12-02 14:35:23 -08:00
Matthew Johnson
768b076420 add an assert primitive 
The assert primitive has an effectful API and so it can't be staged out;
it's only a trace-time primitive. It can be discharged to the functional
form.

We might want to have separate transforms for discharging errors and for
adding error checks. But right now they're just bundled together in the
checkify transform.
 2021-12-02 11:33:56 -08:00
jax authors
7869a6cb75 Merge pull request #8753 from mattjj:checkify 
PiperOrigin-RevId: 413513067
 2021-12-01 14:34:17 -08:00
Matthew Johnson
659f8b794f add skeleton checkify transformation 2021-12-01 10:44:58 -08:00