Modularize remaining unary ops

tfjs-core/src/gradients/Abs_grad.ts

@@ -19,7 +19,7 @@ import {Abs} from '../kernel_names';
 import {GradConfig} from '../kernel_registry';
 import {cast} from '../ops/cast';
 import {mul} from '../ops/mul';
-import {step} from '../ops/unary_ops';
+import {step} from '../ops/step';
 import {Tensor} from '../tensor';
 
 export const absGradConfig: GradConfig = {

tfjs-core/src/gradients/Acos_grad.ts

@@ -21,9 +21,9 @@ import {cast} from '../ops/cast';
 import {div} from '../ops/div';
 import {neg} from '../ops/neg';
 import {scalar} from '../ops/scalar';
+import {sqrt} from '../ops/sqrt';
 import {square} from '../ops/square';
 import {sub} from '../ops/sub';
-import {sqrt} from '../ops/unary_ops';
 import {Tensor} from '../tensor';
 
 export const acosGradConfig: GradConfig = {

tfjs-core/src/gradients/Acosh_grad.ts

@@ -19,9 +19,9 @@ import {Acosh} from '../kernel_names';
 import {GradConfig} from '../kernel_registry';
 import {cast} from '../ops/cast';
 import {div} from '../ops/div';
+import {sqrt} from '../ops/sqrt';
 import {square} from '../ops/square';
 import {sub} from '../ops/sub';
-import {sqrt} from '../ops/unary_ops';
 import {Tensor} from '../tensor';
 
 export const acoshGradConfig: GradConfig = {

tfjs-core/src/gradients/Asin_grad.ts

@@ -20,9 +20,9 @@ import {GradConfig} from '../kernel_registry';
 import {cast} from '../ops/cast';
 import {div} from '../ops/div';
 import {scalar} from '../ops/scalar';
+import {sqrt} from '../ops/sqrt';
 import {square} from '../ops/square';
 import {sub} from '../ops/sub';
-import {sqrt} from '../ops/unary_ops';
 import {Tensor} from '../tensor';
 
 export const asinGradConfig: GradConfig = {

tfjs-core/src/gradients/Asinh_grad.ts

@@ -21,8 +21,8 @@ import {add} from '../ops/add';
 import {cast} from '../ops/cast';
 import {div} from '../ops/div';
 import {scalar} from '../ops/scalar';
+import {sqrt} from '../ops/sqrt';
 import {square} from '../ops/square';
-import {sqrt} from '../ops/unary_ops';
 import {Tensor} from '../tensor';
 
 export const asinhGradConfig: GradConfig = {

tfjs-core/src/gradients/FusedBatchNorm_grad.ts

@@ -20,11 +20,11 @@ import {add} from '../ops/add';
 import {getReductionAxes} from '../ops/broadcast_util';
 import {mul} from '../ops/mul';
 import {reshape} from '../ops/reshape';
+import {rsqrt} from '../ops/rsqrt';
 import {scalar} from '../ops/scalar';
 import {sub} from '../ops/sub';
 import {sum} from '../ops/sum';
 import {tile} from '../ops/tile';
-import {rsqrt} from '../ops/unary_ops';
 import {Tensor} from '../tensor';
 import {Rank, ShapeMap} from '../types';
 

tfjs-core/src/gradients/IsFinite_grad.ts

@@ -0,0 +1,30 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {IsFinite} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {zerosLike} from '../ops/zeros_like';
+import {Tensor} from '../tensor';
+
+export const isFiniteGradConfig: GradConfig = {
+  kernelName: IsFinite,
+  gradFunc: (dy: Tensor) => {
+    // TODO(nsthorat): Let gradients be null for cases where we want to stop
+    // backpropgation.
+    return {x: () => zerosLike(dy)};
+  }
+};

tfjs-core/src/gradients/IsInf_grad.ts

@@ -0,0 +1,31 @@
+
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {IsInf} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {zerosLike} from '../ops/zeros_like';
+import {Tensor} from '../tensor';
+
+export const isInfGradConfig: GradConfig = {
+  kernelName: IsInf,
+  gradFunc: (dy: Tensor) => {
+    // TODO(nsthorat): Let gradients be null for cases where we want to stop
+    // backpropgation.
+    return {x: () => zerosLike(dy)};
+  }
+};

tfjs-core/src/gradients/IsNan_grad.ts

@@ -0,0 +1,30 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {zerosLike} from '../ops/zeros_like';
+import {Tensor} from '../tensor';
+
+export const isNanGradConfig: GradConfig = {
+  kernelName: IsNan,
+  gradFunc: (dy: Tensor) => {
+    // TODO(nsthorat): Let gradients be null for cases where we want to stop
+    // backpropgation.
+    return {x: () => zerosLike(dy)};
+  }
+};

tfjs-core/src/gradients/Relu6_grad.ts

@@ -19,7 +19,7 @@ import {GradConfig} from '../kernel_registry';
 import {cast} from '../ops/cast';
 import {lessEqual} from '../ops/less_equal';
 import {mul} from '../ops/mul';
-import {step} from '../ops/unary_ops';
+import {step} from '../ops/step';
 import {Tensor} from '../tensor';
 
 export const relu6GradConfig: GradConfig = {

tfjs-core/src/gradients/Relu_grad.ts

@@ -18,7 +18,7 @@ import {Relu} from '../kernel_names';
 import {GradConfig} from '../kernel_registry';
 import {cast} from '../ops/cast';
 import {mul} from '../ops/mul';
-import {step} from '../ops/unary_ops';
+import {step} from '../ops/step';
 import {Tensor} from '../tensor';
 
 export const reluGradConfig: GradConfig = {

tfjs-core/src/gradients/Round_grad.ts

@@ -0,0 +1,30 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {Round} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {zerosLike} from '../ops/zeros_like';
+import {Tensor} from '../tensor';
+
+export const roundGradConfig: GradConfig = {
+  kernelName: Round,
+  gradFunc: (dy: Tensor) => {
+    // TODO(nsthorat): Let gradients be null for cases where we want to stop
+    // backpropgation.
+    return {x: () => zerosLike(dy)};
+  }
+};

tfjs-core/src/gradients/Rsqrt_grad.ts

@@ -0,0 +1,33 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {Rsqrt} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {div} from '../ops/div';
+import {mul} from '../ops/mul';
+import {neg} from '../ops/neg';
+import {pow} from '../ops/pow';
+import {Tensor} from '../tensor';
+
+export const rsqrtGradConfig: GradConfig = {
+  kernelName: Rsqrt,
+  inputsToSave: ['x'],
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const [x] = saved;
+    return {x: () => neg(div(dy, mul(pow(x, 1.5), 2)))};
+  }
+};

tfjs-core/src/gradients/Sigmoid_grad.ts

@@ -0,0 +1,33 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {Sigmoid} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {mul} from '../ops/mul';
+import {scalar} from '../ops/scalar';
+import {sub} from '../ops/sub';
+import {Tensor} from '../tensor';
+
+export const sigmoidGradConfig: GradConfig = {
+  kernelName: Sigmoid,
+  outputsToSave: [true],
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const [y] = saved;
+
+    return {x: () => mul(dy, mul(y, sub(scalar(1), y)))};
+  }
+};

tfjs-core/src/gradients/Softplus_grad.ts

@@ -0,0 +1,32 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {Softplus} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {mul} from '../ops/mul';
+import {sigmoid} from '../ops/sigmoid';
+import {Tensor} from '../tensor';
+
+export const softplusGradConfig: GradConfig = {
+  kernelName: Softplus,
+  inputsToSave: ['x'],
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const [x] = saved;
+
+    return {x: () => mul(dy, sigmoid(x))};
+  }
+};

tfjs-core/src/gradients/Sqrt_grad.ts

@@ -0,0 +1,34 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {Sqrt} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {cast} from '../ops/cast';
+import {div} from '../ops/div';
+import {mul} from '../ops/mul';
+import {sqrt} from '../ops/sqrt';
+import {Tensor} from '../tensor';
+
+export const sqrtGradConfig: GradConfig = {
+  kernelName: Sqrt,
+  inputsToSave: ['x'],
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const [x] = saved;
+
+    return {x: () => div(dy, mul(sqrt(cast(x, 'float32')), 2))};
+  }
+};

tfjs-core/src/gradients/Step_grad.ts

@@ -0,0 +1,30 @@
+/**
+ * @license
+ * Copyright 2020 Google LLC. 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 {Step} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {zerosLike} from '../ops/zeros_like';
+import {Tensor} from '../tensor';
+
+export const stepGradConfig: GradConfig = {
+  kernelName: Step,
+  gradFunc: (dy: Tensor) => {
+    // TODO(manrajgrover): Return null for gradients when backprop supports
+    // it.
+    return {x: () => zerosLike(dy)};
+  }
+};