Function MulNoNanImpl

tensorflow/compiler/tf2xla/kernels/binary_ops.cc:91–99

Implementation of MulNoNan. Pseudo-code: if (y == 0) { return 0 } else { return x * y; }

//   return x * y;
// }
static xla::XlaOp MulNoNanImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  auto zero = XlaHelpers::Zero(b, dtype);
  auto y_equals_0 = xla::Eq(y, zero);
  auto zeros = xla::ZerosLike(x);
  auto result = xla::Select(y_equals_0, zeros, xla::Mul(x, y));
  return result;
}
XLA_MAKE_BINARY(MulNoNan,
                MulNoNanImpl(b, input_type(0), lhs, rhs, broadcast_helper));