experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add a CPU bfloat16 kernel for MulNoNan

Browse Source 
Fixes a test whose randomness was affected by this.

PiperOrigin-RevId: 309495955
Change-Id: I4e184c9ca40b0102cf5eb895a040f9f2a09e753a
This commit is contained in:
Allen Lavoie 2020-05-01 16:34:03 -07:00 committed by TensorFlower Gardener
parent c161d14499
commit be894c5c42
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/compiler/tests/ternary_ops_test.py
View File

@ -230,7 +230,7 @@ class TernaryOpsTest(xla_test.XLATestCase, parameterized.TestCase):
{
'sigma': 1e15,
'rtol': 1e-6,
'atol': 1e-6
'atol': 1e-4
},
{
'sigma': 30,
@ -240,7 +240,7 @@ class TernaryOpsTest(xla_test.XLATestCase, parameterized.TestCase):
{
'sigma': 1e-8,
'rtol': 5e-4,
'atol': 3e-6
'atol': 3e-4
},
{
'sigma': 1e-16,

4
tensorflow/core/kernels/cwise_op_mul_1.cc
View File

@ -19,8 +19,8 @@ namespace tensorflow {
REGISTER6(BinaryOp, CPU, "Mul", functor::mul, float, Eigen::half, double, uint8,
int32, bfloat16);
REGISTER5(BinaryOp, CPU, "MulNoNan", functor::mul_no_nan, Eigen::half, float,
double, complex64, complex128);
REGISTER6(BinaryOp, CPU, "MulNoNan", functor::mul_no_nan, Eigen::half, float,
double, complex64, complex128, bfloat16);
#if defined(__ANDROID_TYPES_SLIM__)
// We only register the first type when we have multi-argument calls in the