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

disable for now

Browse Source
This commit is contained in:
Matthew Johnson 2022-04-29 11:03:15 -07:00
parent 36b0506b4d
commit b4a4ffd381
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/lax/control_flow.py
View File

@ -2079,7 +2079,8 @@ core.custom_typechecks[scan_p] = partial(_scan_typecheck, False)
pe.partial_eval_jaxpr_custom_rules[scan_p] = \
partial(pe.partial_eval_jaxpr_custom_rule_not_implemented, 'scan')
pe.padding_rules[scan_p] = _scan_padding_rule
pe.dce_rules[scan_p] = _scan_dce_rule
# TODO(mattjj): re-enable
# pe.dce_rules[scan_p] = _scan_dce_rule
@api_boundary