mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix typo

Browse Source
This commit is contained in:
Jacob Kelly 2020-02-28 20:40:47 -05:00 committed by Matthew Johnson
parent ddd52c4730
commit 098aabefcd
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/lax/lax.py
View File

@ -3184,6 +3184,7 @@ gather_p = standard_primitive(
_gather_shape_rule, _gather_dtype_rule, 'gather',
_gather_translation_rule)
ad.defjvp(gather_p, _gather_jvp_rule, None)
ad.primitive_transposes[gather_p] = _gather_transpose_rule
batching.primitive_batchers[gather_p] = _gather_batching_rule