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 brcast to deal with inconsistent shapes.

Browse Source
This commit is contained in:
Yutong Zhao 2019-04-17 20:54:01 -04:00
parent dfd3d93350
commit bbf0d5c55e
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/lax/lax.py
View File

@ -1442,8 +1442,8 @@ ad.defjvp(cos_p, lambda g, x: neg(mul(g, sin(x))))
atan2_p = standard_binop([_float, _float], 'atan2')
ad.defjvp(atan2_p,
lambda g, x, y: g * y / (square(x) + square(y)),
lambda g, x, y: g * -x / (square(x) + square(y)))
lambda g, x, y: _brcast(g, y) * (y / (square(x) + square(y))),
lambda g, x, y: _brcast(g, x) * -x / (square(x) + square(y)))
lgamma_p = standard_unop(_float, 'lgamma')
ad.defjvp(lgamma_p, lambda g, x: mul(g, digamma(x)))