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

[Mosaic TPU] Fix operands order in try canonicalize add of matmul.

Browse Source 
PiperOrigin-RevId: 671437100
This commit is contained in:
Jevin Jiang 2024-09-05 11:05:54 -07:00 committed by jax authors
parent 97db78ba24
commit dba674153e
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jaxlib/mosaic/dialect/tpu/tpu_ops.cc
View File

@ -493,8 +493,9 @@ class CanonicalizeAddOfMatmul : public OpRewritePattern<AddOp> {
}
return failure();
};
return success(succeeded(try_canonicalize(op.getLhs(), op.getRhs())) ||
succeeded(try_canonicalize(op.getLhs(), op.getRhs())));
// We tried try_canonicalize(op.getRhs(), op.getLhs()) and it caused
// worrying numerical differences in some of kernels.
return try_canonicalize(op.getLhs(), op.getRhs());
}
};