[wasm] implement isNaN for wasm

tfjs-backend-wasm/src/cc/BUILD

@@ -323,6 +323,7 @@ tfjs_cc_library(
         ":GatherNd",
         ":Greater",
         ":GreaterEqual",
+        ":IsNan",
         ":LeakyRelu",
         ":Less",
         ":LessEqual",
@@ -1299,3 +1300,12 @@ tfjs_cc_library(
         ":unary",
     ],
 )
+
+tfjs_cc_library(
+    name = "IsNan",
+    srcs = ["kernels/IsNan.cc"],
+    deps = [
+        ":backend",
+        ":unary",
+    ],
+)

tfjs-backend-wasm/src/cc/kernels/IsNan.cc

@@ -0,0 +1,38 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * ===========================================================================*/
+
+#ifdef __EMSCRIPTEN__
+#include <emscripten.h>
+#endif
+
+#include <math.h>
+
+#include "tfjs-backend-wasm/src/cc/backend.h"
+#include "tfjs-backend-wasm/src/cc/unary.h"
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void IsNan(const int x_id, const DType dtype, const int out_id) {
+  unary_f32_with_bool_out(x_id, out_id, isnan);
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

tfjs-backend-wasm/src/cc/unary.h

@@ -24,8 +24,7 @@
 namespace tfjs {
 namespace wasm {
 
-inline void unary_bool(const int x_id, const int out_id,
-                       bool operation(bool)) {
+inline void unary_bool(const int x_id, const int out_id, bool operation(bool)) {
   auto& a_info = backend::get_tensor_info(x_id);
   auto& out_info = backend::get_tensor_info_out(out_id);
 
@@ -76,6 +75,19 @@ inline void unary_i32_with_f32_out(const size_t x_id, const size_t out_id,
   }
 }
 
+inline void unary_f32_with_bool_out(const size_t x_id, const size_t out_id,
+                                    bool operation(float)) {
+  auto& a_info = backend::get_tensor_info(x_id);
+  auto& out_info = backend::get_tensor_info_out(out_id);
+
+  const float* a_buf = a_info.f32();
+  bool* out_buf = out_info.b_write();
+
+  for (size_t i = 0; i < a_info.size; ++i) {
+    out_buf[i] = operation(a_buf[i]);
+  }
+}
+
 typedef xnn_status (*xnn_create_unary_op)(size_t, size_t, size_t, uint32_t,
                                           xnn_operator_t*);
 typedef xnn_status (*xnn_setup_unary_op)(xnn_operator_t, size_t, const float*,

tfjs-backend-wasm/src/kernels/IsNan.ts

@@ -0,0 +1,22 @@
+/**
+ * @license
+ * Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the License);
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an AS IS BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {IsNan, KernelConfig} from '@tensorflow/tfjs-core';
+
+import {createUnaryKernelConfig} from './unary_kernel';
+
+export const isNaNConfig: KernelConfig = createUnaryKernelConfig(IsNan, 'bool');

tfjs-backend-wasm/src/register_all_kernels.ts

@@ -58,6 +58,7 @@ import {gatherV2Config} from './kernels/GatherV2';
 import {greaterConfig} from './kernels/Greater';
 import {greaterEqualConfig} from './kernels/GreaterEqual';
 import {identityConfig} from './kernels/Identity';
+import {isNaNConfig} from './kernels/IsNan';
 import {leakyReluConfig} from './kernels/LeakyRelu';
 import {lessConfig} from './kernels/Less';
 import {lessEqualConfig} from './kernels/LessEqual';
@@ -170,6 +171,7 @@ const kernelConfigs: KernelConfig[] = [
   greaterConfig,
   greaterEqualConfig,
   identityConfig,
+  isNaNConfig,
   leakyReluConfig,
   lessConfig,
   lessEqualConfig,

tfjs-backend-wasm/src/setup_test.ts

@@ -399,7 +399,8 @@ const TEST_FILTERS: TestFilter[] = [
   {include: 'stringNGrams'},
   {include: 'stringSplit'},
   {include: 'stringToHashBucketFast'},
-  {startsWith: 'reciprocal'},
+  {include: 'reciprocal'},
+  {include: 'isNaN'},
 ];
 
 const customInclude = (testName: string) => {