add modularized gradient for squared difference

tafsiri · tafsiri · commit d21efb03d6bf · 2020-02-19T17:30:37.000-05:00
diff --git a/tfjs-core/src/backends/cpu/all_kernels.ts b/tfjs-core/src/backends/cpu/all_kernels.ts
@@ -21,7 +21,7 @@
 import './non_max_suppression_v5';
 import './square';
 
-// TODO this will be a dist import from types to avoid a circular dependency
+// TODO this will be a dist import from types
 import {KernelConfig} from '../../kernel_registry';
 
 // Import Kernel Configs here.
diff --git a/tfjs-core/src/gradients/SquaredDifference_grad.ts b/tfjs-core/src/gradients/SquaredDifference_grad.ts
@@ -0,0 +1,33 @@
+/**
+ * @license
+ * Copyright 2010 Google Inc. 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 {SquaredDifference} from '../kernel_names';
+import {GradConfig} from '../kernel_registry';
+import {mul, sub} from '../ops/binary_ops';
+import {scalar} from '../ops/tensor_ops';
+import {Tensor} from '../tensor';
+
+export const squaredDifferenceGrad_: GradConfig = {
+  kernelName: SquaredDifference,
+  gradFunc: (dy: Tensor, saved: Tensor[]) => {
+    const [$a, $b] = saved;
+    const two = scalar(2);
+    const derA = () => mul(dy, mul(two, sub($a, $b)));
+    const derB = () => mul(dy, mul(two, sub($b, $a)));
+    return {$a: derA, $b: derB};
+  }
+};
diff --git a/tfjs-core/src/index.ts b/tfjs-core/src/index.ts
@@ -37,6 +37,10 @@ import './backends/cpu/backend_cpu';
 import './backends/cpu/register_all_kernels';
 // Import all kernels from webgl.
 import './backends/webgl/register_all_kernels';
+
+// Register all the gradients
+import './register_all_gradients';
+
 import './platforms/platform_browser';
 import './platforms/platform_node';
 
diff --git a/tfjs-core/src/ops/all_gradients.ts b/tfjs-core/src/ops/all_gradients.ts
@@ -20,3 +20,12 @@
 // the contents of this file and import only the gradients that are needed.
 
 import './square_grad';
+
+// Import Grad Configs here.
+import {squaredDifferenceGrad_} from '../gradients/SquaredDifference_grad';
+import {GradConfig} from '../kernel_registry';
+
+// Export all kernel configs here so that the package can auto register them
+export const gradConfigs: GradConfig[] = [
+  squaredDifferenceGrad_,
+];
diff --git a/tfjs-core/src/ops/ops.ts b/tfjs-core/src/ops/ops.ts
@@ -15,9 +15,6 @@
  * =============================================================================
  */
 
-// Importing this file registers gradients in the global registry.
-import './all_gradients';
-
 // Modularized ops.
 export {square} from './square';
 export {squaredDifference} from './squared_difference';
diff --git a/tfjs-core/src/register_all_gradients.ts b/tfjs-core/src/register_all_gradients.ts
@@ -0,0 +1,22 @@
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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 {registerGradient} from './kernel_registry';
+import {gradConfigs} from './ops/all_gradients';
+
+for (const gradientConfig of gradConfigs) {
+  registerGradient(gradientConfig);
+}
