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 scalar broadcast bug in safe_mul translation

Browse Source
This commit is contained in:
Matthew Johnson 2019-03-23 17:10:01 -07:00
parent 87143d8e0a
commit b041435f29
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.py
View File

@ -1732,8 +1732,8 @@ ad.defbilinear_broadcasting(_brcast, mul_p, mul, mul) # TODO
def _safe_mul_translation_rule(c, x, y):
dtype = c.GetShape(x).numpy_dtype()
zero = c.Constant(onp.array(0, dtype=dtype))
out_shape = tuple(onp.maximum(c.GetShape(x).dimensions(),
c.GetShape(y).dimensions()))
out_shape = broadcast_shapes(c.GetShape(x).dimensions(),
c.GetShape(y).dimensions())
return c.Select(c.Or(c.Eq(x, zero), c.Eq(y, zero)),
c.Broadcast(zero, out_shape),
c.Mul(x, y))