modularize reverse ops

tfjs-core/src/gradients/Reverse_grad.ts

@@ -0,0 +1,30 @@
+/**
+ * @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 {Reverse, ReverseAttrs} from '../kernel_names';
+import {GradConfig, NamedAttrMap} from '../kernel_registry';
+import {reverse} from '../ops/reverse';
+import {Tensor} from '../tensor';
+import {parseAxisParam} from '../util';
+
+export const reverseGradConfig: GradConfig = {
+  kernelName: Reverse,
+  gradFunc: (dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => {
+    const {dims} = attrs as {} as ReverseAttrs;
+    const axes = parseAxisParam(dims, dy.shape);
+    return {x: () => reverse(dy, axes)};
+  }
+};

tfjs-core/src/kernel_names.ts

@@ -440,6 +440,12 @@ export type ResizeBilinearGradInputs = Pick<NamedTensorInfoMap, 'images'>;
 export const Relu6 = 'Relu6';
 export type Relu6Inputs = Pick<NamedTensorInfoMap, 'x'>;
 
+export const Reverse = 'Reverse';
+export type ReverseInputs = Pick<NamedTensorInfoMap, 'x'>;
+export interface ReverseAttrs {
+  dims: number|number[];
+}
+
 export const SelectV2 = 'SelectV2';
 export type SelectV2Inputs = Pick<NamedTensorInfoMap, 'condition'|'t'|'e'>;
 

tfjs-core/src/ops/ops.ts

@@ -93,6 +93,11 @@ export {randomUniform} from './random_uniform';
 export {real} from './real';
 export {relu} from './relu';
 export {relu6} from './relu6';
+export {reverse} from './reverse';
+export {reverse1d} from './reverse_1d';
+export {reverse2d} from './reverse_2d';
+export {reverse3d} from './reverse_3d';
+export {reverse4d} from './reverse_4d';
 export {selu} from './selu';
 export {separableConv2d} from './separable_conv2d';
 export {spaceToBatchND} from './space_to_batch_nd';
@@ -106,7 +111,6 @@ export {where} from './where';
 export {whereAsync} from './where_async';
 
 export * from './boolean_mask';
-export * from './reverse';
 export * from './slice';
 export * from './unary_ops';
 export * from './reduction_ops';

tfjs-core/src/ops/reverse.ts

@@ -15,70 +15,18 @@
  * =============================================================================
  */
 
-import {ENGINE} from '../engine';
-import {Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D} from '../tensor';
+import {ENGINE, ForwardFunc} from '../engine';
+import {Reverse, ReverseAttrs, ReverseInputs} from '../kernel_names';
+import {NamedAttrMap} from '../kernel_registry';
+import {Tensor} from '../tensor';
+import {NamedTensorMap} from '../tensor_types';
 import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
-import * as util from '../util';
-import {op} from './operation';
+import {parseAxisParam} from '../util';
 
-/**
- * Reverses a `tf.Tensor1D`.
- *
- * @param x The input tensor.
- */
-function reverse1d_(x: Tensor1D|TensorLike): Tensor1D {
-  const $x = convertToTensor(x, 'x', 'reverse');
-  util.assert(
-      $x.rank === 1,
-      () => `Error in reverse1D: x must be rank 1 but got rank ${$x.rank}.`);
-  return reverse($x, 0);
-}
-
-/**
- * Reverses a `tf.Tensor2D` along a specified axis.
- *
- * @param x The input tensor.
- * @param axis The set of dimensions to reverse. Must be in the
- *     range [-rank(x), rank(x)). Defaults to all axes.
- */
-function reverse2d_(x: Tensor2D|TensorLike, axis?: number|number[]): Tensor2D {
-  const $x = convertToTensor(x, 'x', 'reverse');
-  util.assert(
-      $x.rank === 2,
-      () => `Error in reverse2D: x must be rank 2 but got rank ${$x.rank}.`);
-  return reverse($x, axis);
-}
-
-/**
- * Reverses a `tf.Tensor3D` along a specified axis.
- *
- * @param x The input tensor.
- * @param axis The set of dimensions to reverse. Must be in the
- *     range [-rank(x), rank(x)). Defaults to all axes.
- */
-function reverse3d_(x: Tensor3D|TensorLike, axis?: number|number[]): Tensor3D {
-  const $x = convertToTensor(x, 'x', 'reverse');
-  util.assert(
-      $x.rank === 3,
-      () => `Error in reverse3D: x must be rank 3 but got rank ${$x.rank}.`);
-  return reverse($x, axis);
-}
-
-/**
- * Reverses a `tf.Tensor4D` along a specified axis.
- *
- * @param x The input tensor.
- * @param axis The set of dimensions to reverse. Must be in the
- *     range [-rank(x), rank(x)). Defaults to all axes.
- */
-function reverse4d_(x: Tensor4D|TensorLike, axis?: number|number[]): Tensor4D {
-  const $x = convertToTensor(x, 'x', 'reverse');
-  util.assert(
-      $x.rank === 4,
-      () => `Error in reverse4D: x must be rank 4 but got rank ${$x.rank}.`);
-  return reverse($x, axis);
-}
+import {reshape} from './array_ops';
+import {clone} from './clone';
+import {op} from './operation';
 
 /**
  * Reverses a `tf.Tensor` along a specified axis.
@@ -114,20 +62,21 @@ function reverse_<T extends Tensor>(
     x: T|TensorLike, axis?: number|number[]): T {
   const $x = convertToTensor(x, 'x', 'reverse');
 
-  if ($x.rank === 0) {
-    return $x.clone();
-  }
-  const axes = util.parseAxisParam(axis, $x.shape);
-  const grad = (dy: T) => {
-    return {$x: () => dy.reverse(axes)};
+  const forward: ForwardFunc<Tensor> = (backend) => {
+    const axes = parseAxisParam(axis, $x.shape);
+    if ($x.rank === 0) {
+      return clone($x);
+    }
+    const res = backend.reverse($x, axes);
+    return reshape(res, $x.shape);
   };
-  const res =
-      ENGINE.runKernelFunc(backend => backend.reverse($x, axes), {$x}, grad);
-  return res.reshapeAs($x);
+
+  const inputs: ReverseInputs = {x: $x};
+  const attrs: ReverseAttrs = {dims: axis};
+
+  return ENGINE.runKernelFunc(
+             forward, inputs as {} as NamedTensorMap, null /* gradient */,
+             Reverse, attrs as {} as NamedAttrMap) as T;
 }
 
 export const reverse = op({reverse_});
-export const reverse1d = op({reverse1d_});
-export const reverse2d = op({reverse2d_});
-export const reverse3d = op({reverse3d_});
-export const reverse4d = op({reverse4d_});

tfjs-core/src/ops/reverse_1d.ts

@@ -0,0 +1,38 @@
+/**
+ * @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 {Tensor1D} from '../tensor';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+import * as util from '../util';
+import {op} from './operation';
+import {reverse} from './reverse';
+
+/**
+ * Reverses a `tf.Tensor1D`.
+ *
+ * @param x The input tensor.
+ */
+function reverse1d_(x: Tensor1D|TensorLike): Tensor1D {
+  const $x = convertToTensor(x, 'x', 'reverse');
+  util.assert(
+      $x.rank === 1,
+      () => `Error in reverse1D: x must be rank 1 but got rank ${$x.rank}.`);
+  return reverse($x, 0);
+}
+
+export const reverse1d = op({reverse1d_});

tfjs-core/src/ops/reverse_1d_test.ts

@@ -0,0 +1,60 @@
+/**
+ * @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 * as tf from '../index';
+import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
+import {expectArraysClose} from '../test_util';
+
+describeWithFlags('reverse1d', ALL_ENVS, () => {
+  it('reverse a 1D array', async () => {
+    const input = tf.tensor1d([1, 2, 3, 4, 5]);
+    const result = tf.reverse1d(input);
+    expect(result.shape).toEqual(input.shape);
+    expectArraysClose(await result.data(), [5, 4, 3, 2, 1]);
+  });
+
+  it('reverse a 1D array, even length', async () => {
+    const input = tf.tensor1d([1, 2, 3, 4]);
+    const result = tf.reverse1d(input);
+    expect(result.shape).toEqual(input.shape);
+    expectArraysClose(await result.data(), [4, 3, 2, 1]);
+  });
+
+  it('grad', async () => {
+    const a = tf.tensor1d([1, 2, 3]);
+    const dy = tf.tensor1d([10, 20, 30]);
+    const da = tf.grad((a: tf.Tensor1D) => tf.reverse1d(a))(a, dy);
+    expect(da.shape).toEqual([3]);
+    expectArraysClose(await da.data(), [30, 20, 10]);
+  });
+
+  it('gradient with clones', async () => {
+    const a = tf.tensor1d([1, 2, 3]);
+    const dy = tf.tensor1d([10, 20, 30]);
+    const da =
+        tf.grad((a: tf.Tensor1D) => tf.reverse1d(a.clone()).clone())(a, dy);
+    expect(da.shape).toEqual([3]);
+    expectArraysClose(await da.data(), [30, 20, 10]);
+  });
+
+  it('accepts a tensor-like object', async () => {
+    const input = [1, 2, 3, 4, 5];
+    const result = tf.reverse1d(input);
+    expect(result.shape).toEqual([5]);
+    expectArraysClose(await result.data(), [5, 4, 3, 2, 1]);
+  });
+});

tfjs-core/src/ops/reverse_2d.ts

@@ -0,0 +1,40 @@
+/**
+ * @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 {Tensor2D} from '../tensor';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+import * as util from '../util';
+import {op} from './operation';
+import {reverse} from './reverse';
+
+/**
+ * Reverses a `tf.Tensor2D` along a specified axis.
+ *
+ * @param x The input tensor.
+ * @param axis The set of dimensions to reverse. Must be in the
+ *     range [-rank(x), rank(x)). Defaults to all axes.
+ */
+function reverse2d_(x: Tensor2D|TensorLike, axis?: number|number[]): Tensor2D {
+  const $x = convertToTensor(x, 'x', 'reverse');
+  util.assert(
+      $x.rank === 2,
+      () => `Error in reverse2D: x must be rank 2 but got rank ${$x.rank}.`);
+  return reverse($x, axis);
+}
+
+export const reverse2d = op({reverse2d_});