fix psum transpose rule

jax-ml · May 1, 2020 · 268e0d1 · 268e0d1
1 parent 564e4a2
commit 268e0d1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
@@ -341,7 +341,7 @@ def _translate(val):
     partial(_allreduce_split_axis_rule, psum_p, lax._reduce_sum)
 xla.parallel_translations[psum_p] = _psum_translation_rule
 pxla.parallel_pure_rules[psum_p] = lambda *args, shape: (x * prod(shape) for x in args)
-ad.deflinear(psum_p, lambda *ts, axis_name, axis_index_groups: psum(
+ad.deflinear(psum_p, lambda ts, axis_name, axis_index_groups: psum_p.bind(
     *ts, axis_name=axis_name, axis_index_groups=axis_index_groups))
 pxla.multi_host_supported_collectives.add(psum_p)