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

Add defensive tuple() in lax.reduce_window (#3741)

Browse Source
This commit is contained in:
James Bradbury 2020-07-13 14:37:46 -07:00 committed by GitHub
parent 3c6cd5fb94
commit 6017205cea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/lax/lax.py
View File

@ -1123,6 +1123,8 @@ def reduce_window(operand: Array, init_value: Array, computation: Callable,
if isinstance(padding, str):
padding = padtype_to_pads(operand.shape, window_dimensions,
window_strides, padding)
else:
padding = tuple(padding)
monoid_reducer = _get_monoid_window_reducer(computation, init_value)
if monoid_reducer:
return monoid_reducer(operand, window_dimensions, window_strides, padding)