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

Fix minor issue in conj_p translation rule.

Browse Source 
conj_p was forwarding its keyword arguments to the ComputationBuilder.Conj() method. The current implementation of Conj() ignores extra keyword arguments, but we shouldn't depend on this kind of implementation detail.
This commit is contained in:
Peter Hawkins 2019-03-26 10:31:30 -04:00
parent c9aa60102f
commit fa6b06fd15
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.py
View File

@ -1669,6 +1669,7 @@ def _conj_transpose_rule(t, x, input_dtype):
else:
return [real(t)]
xla.translations[conj_p] = lambda c, x, **kwargs: c.Conj(x)
ad.primitive_jvps[conj_p] = partial(ad.linear_jvp, conj_p)
ad.primitive_transposes[conj_p] = _conj_transpose_rule