[wasm] Add acos and acosh kernels

tfjs-backend-wasm/src/cc/BUILD.bazel

@@ -293,6 +293,8 @@ tfjs_cc_library(
     name = "all_kernels",
     deps = [
         ":Abs",
+        ":Acos",
+        ":Acosh",
         ":Add",
         ":AddN",
         ":All",
@@ -383,6 +385,26 @@ tfjs_cc_library(
     ],
 )
 
+tfjs_cc_library(
+    name = "Acos",
+    srcs = ["kernels/Acos.cc"],
+    deps = [
+        ":backend",
+        ":unary",
+        ":util",
+    ],
+)
+
+tfjs_cc_library(
+    name = "Acosh",
+    srcs = ["kernels/Acosh.cc"],
+    deps = [
+        ":backend",
+        ":unary",
+        ":util",
+    ],
+)
+
 tfjs_cc_library(
     name = "Add",
     srcs = ["kernels/Add.cc"],

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

@@ -0,0 +1,61 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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 <cmath>
+
+#include "tfjs-backend-wasm/src/cc/backend.h"
+#include "tfjs-backend-wasm/src/cc/unary.h"
+#include "tfjs-backend-wasm/src/cc/util.h"
+
+namespace {
+
+template <typename T>
+inline T acos_impl(T n) {
+  return static_cast<T>(std::acosf(static_cast<float>(n)));
+}
+
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void Acos(const int x_id, const DType dtype, const int out_id) {
+  switch (dtype) {
+    case DType::float32:
+      unary_f32(x_id, out_id, acos_impl<float>);
+      break;
+    case DType::int32:
+      unary_i32(x_id, out_id, acos_impl<int>);
+      break;
+    default:
+      util::warn("Acos for tensor id %d failed. Unsupported dtype %d", x_id,
+                 dtype);
+  }
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

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

@@ -0,0 +1,58 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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 <cmath>
+
+#include "tfjs-backend-wasm/src/cc/backend.h"
+#include "tfjs-backend-wasm/src/cc/unary.h"
+#include "tfjs-backend-wasm/src/cc/util.h"
+
+namespace {
+
+template <typename T>
+inline T acosh_impl(T n) {
+  return static_cast<T>(std::acoshf(static_cast<float>(n)));
+}
+
+}  // namespace
+
+namespace tfjs {
+namespace wasm {
+// We use C-style API to interface with Javascript.
+extern "C" {
+
+#ifdef __EMSCRIPTEN__
+EMSCRIPTEN_KEEPALIVE
+#endif
+void Acosh(const int x_id, const DType dtype, const int out_id) {
+  switch (dtype) {
+    case DType::float32:
+      unary_f32(x_id, out_id, acosh_impl<float>);
+      break;
+    default:
+      util::warn("Acosh for tensor id %d failed. Unsupported dtype %d", x_id,
+                 dtype);
+  }
+}
+
+}  // extern "C"
+}  // namespace wasm
+}  // namespace tfjs

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

@@ -28,8 +28,8 @@
 namespace {
 
 template <typename T>
-inline float atan_impl(T n) {
-  return std::atanf(static_cast<float>(n));
+inline T atan_impl(T n) {
+  return static_cast<T>(std::atanf(static_cast<float>(n)));
 }
 
 }  // namespace
@@ -48,7 +48,7 @@ void Atan(const int x_id, const DType dtype, const int out_id) {
       unary_f32(x_id, out_id, atan_impl<float>);
       break;
     case DType::int32:
-      unary_i32_with_f32_out(x_id, out_id, atan_impl<int>);
+      unary_i32(x_id, out_id, atan_impl<int>);
       break;
     default:
       util::warn("Atan for tensor id %d failed. Unsupported dtype %d", x_id,

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

@@ -0,0 +1,22 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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 {Acos, KernelConfig} from '@tensorflow/tfjs-core';
+
+import {createUnaryKernelConfig} from './unary_kernel';
+
+export const acosConfig: KernelConfig = createUnaryKernelConfig(Acos);

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

@@ -0,0 +1,22 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC.
+ * 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 {Acosh, KernelConfig} from '@tensorflow/tfjs-core';
+
+import {createUnaryKernelConfig} from './unary_kernel';
+
+export const acoshConfig: KernelConfig = createUnaryKernelConfig(Acosh);

tfjs-backend-wasm/src/register_all_kernels.ts

@@ -21,6 +21,8 @@ import {KernelConfig, registerKernel} from '@tensorflow/tfjs-core';
 
 import {_fusedMatMulConfig} from './kernels/_FusedMatMul';
 import {absConfig} from './kernels/Abs';
+import {acosConfig} from './kernels/Acos';
+import {acoshConfig} from './kernels/Acosh';
 import {addConfig} from './kernels/Add';
 import {addNConfig} from './kernels/AddN';
 import {allConfig} from './kernels/All';
@@ -136,6 +138,8 @@ import {zerosLikeConfig} from './kernels/ZerosLike';
 const kernelConfigs: KernelConfig[] = [
   _fusedMatMulConfig,
   absConfig,
+  acosConfig,
+  acoshConfig,
   addConfig,
   addNConfig,
   allConfig,

tfjs-backend-wasm/src/setup_test.ts

@@ -403,6 +403,8 @@ const TEST_FILTERS: TestFilter[] = [
   {include: 'reciprocal'},
   {include: 'isNaN'},
   {include: 'atan '},
+  {include: 'acos '},
+  {include: 'acosh '},
 ];
 
 const customInclude = (testName: string) => {