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

51 Commits

Author SHA1 Message Date
lenamartens
5ebd81f573 Fix scan and name_stack, rewrite cond to use jaxpr_to_checkify_jaxpr. 
Fix map primitives, test pmap more.
 2023-01-05 16:35:40 +00:00
lenamartens
0bce1cf129 Checkify: switch to initial-style. 2023-01-05 16:35:02 +00:00
lenamartens
3134797968 Add checkify.debug_check which is a noop outside of checkify. 2022-12-14 11:15:34 +00:00
Jake VanderPlas
f09fd8a4e9 [x64] minor test-only updates for better type safety 2022-11-30 15:18:40 -08:00
lenamartens
e4757e8410 Rewrite Checkify to support tracking different error types. 
In general, behavior should remain the same and this is not a breaking
change.

There are some minor changes to the API:
  - checkify.ErrorCategory has changed type: it's no longer an Enum, but
    the JaxException type. These have not been exposed as part of the
    public API.
  - some attributes on Error have changed and made private
  - The raised error has changed type (JaxRuntimeError), and will have a
    different traceback (pointing to the origin of the error + where the
    error value was raised).
  - `checkify.check` now supports formating error message with variable
    size runtime info!

Co-authored-by: Sharad Vikram <sharad.vikram@gmail.com>
 2022-11-25 15:31:54 +00:00
Lena Martens
3116ed52a9 Checkify: fix nan check when primitive has multiple results. 
PiperOrigin-RevId: 488653856
 2022-11-15 07:35:54 -08:00
Yash Katariya
c42bad85ef Make MeshPspecSharding an alias for NamedSharding (it was the other way around before this CL). 
PiperOrigin-RevId: 488473538
 2022-11-14 14:44:00 -08:00
lenamartens
053b8b5bcd Checkify: fix nan_checks+PRNGKeys - a PRNGKey is never NaN! 
Add a guard to the nan_error_rule to not call jnp.isnan on keys.
 2022-11-09 17:08:21 +00:00
lenamartens
c2a00a0526 Disallow checkify-of-vmap-of-while. 2022-10-17 23:01:43 +01:00
Peter Hawkins
0d3277b5c3 Port more tests from jtu.cases_from_list to jtu.sample_product. 2022-10-11 21:06:08 +00:00
lenamartens
0639aced5b Raise cond index into tracing context in case of effects. 
So even if the cond is not data dependent at all, it's included in the
dynamic trace, and effects can be discharged.
 2022-09-29 11:36:04 +01:00
Yash Katariya
9e4114f0f1 Move array.py and sharding.py from experimental/ to _src/. 
PiperOrigin-RevId: 477201711
 2022-09-27 10:06:52 -07:00
lenamartens
27e3981d52 lowerable errors behind a config flag. 2022-09-26 17:34:27 +01:00
jax authors
9c66569514 Merge pull request #12468 from LenaMartens:checkify-but-better 
PiperOrigin-RevId: 476901601
 2022-09-26 08:23:02 -07:00
lenamartens
7078f81dd0 Checkify: misc improvements. 
- err.throw == check_error(err) -> meaning they have the same behavior
  under checkify now
- "divided by zero" -> "division by zero"
- add validation that check_error only takes args of type Error
 2022-09-23 14:33:06 +01:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". 
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
 2022-09-22 12:27:19 -07:00
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