Modularize softmax/logSoftmax, sparseToDense, TopK, stridedSlice

tfjs-core/src/gradients/LogSoftmax_grad.ts

@@ -0,0 +1,37 @@
+/**
+ * @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 {LogSoftmax, LogSoftmaxAttrs} from '../kernel_names';
+import {GradConfig, NamedAttrMap} from '../kernel_registry';
+import {Tensor} from '../tensor';
+
+export const logSoftmaxGradConfig: GradConfig = {
+  kernelName: LogSoftmax,
+  inputsToSave: [],
+  outputsToSave: [true],
+  gradFunc: (dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => {
+    const [value] = saved;
+    const {axis} = attrs as {} as LogSoftmaxAttrs;
+    return {
+      logits: () => {
+        const keepDims = true;
+        const softmax = value.exp();
+        return dy.sub(dy.sum(axis, keepDims).mul(softmax));
+      }
+    };
+  }
+};

tfjs-core/src/gradients/Softmax_grad.ts

@@ -0,0 +1,38 @@
+/**
+ * @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 {Softmax, SoftmaxAttrs} from '../kernel_names';
+import {GradConfig, NamedAttrMap} from '../kernel_registry';
+import {mul} from '../ops/mul';
+import {sub} from '../ops/sub';
+import {sum} from '../ops/sum';
+import {Tensor} from '../tensor';
+
+export const softmaxGradConfig: GradConfig = {
+  kernelName: Softmax,
+  outputsToSave: [true],
+  gradFunc: (dy: Tensor, saved: Tensor[], attrs: NamedAttrMap) => {
+    const [y] = saved;
+    const {dim} = attrs as {} as SoftmaxAttrs;
+    const keepDims = true;
+
+    const dyTimesY = mul(dy, y);
+    return {
+      logits: () => sub(dyTimesY, mul(sum(dyTimesY, [dim], keepDims), y))
+    };
+  }
+};

tfjs-core/src/kernel_names.ts

@@ -376,6 +376,12 @@ export type LogicalNotInputs = Pick<NamedTensorInfoMap, 'x'>;
 export const LogicalOr = 'LogicalOr';
 export type LogicalOrInputs = BinaryInputs;
 
+export const LogSoftmax = 'LogSoftmax';
+export type LogSoftmaxInputs = Pick<NamedTensorInfoMap, 'logits'>;
+export interface LogSoftmaxAttrs {
+  axis: number;
+}
+
 export const LRN = 'LRN';
 export type LRNInputs = Pick<NamedTensorInfoMap, 'x'>;
 export interface LRNAttrs {
@@ -637,6 +643,12 @@ export interface SplitVAttrs {
   axis: number;
 }
 
+export const Softmax = 'Softmax';
+export type SoftmaxInputs = Pick<NamedTensorInfoMap, 'logits'>;
+export interface SoftmaxAttrs {
+  dim: number;
+}
+
 export const SquaredDifference = 'SquaredDifference';
 export type SquaredDifferenceInputs = BinaryInputs;
 
@@ -646,6 +658,26 @@ export type SquareInputs = Pick<NamedTensorInfoMap, 'x'>;
 export const Sub = 'Sub';
 export type SubInputs = BinaryInputs;
 
+export const SparseToDense = 'SparseToDense';
+export type SparseToDenseInputs =
+    Pick<NamedTensorInfoMap, 'sparseIndices'|'sparseValues'|'defaultValue'>;
+export interface SparseToDenseAttrs {
+  outputShape: number[];
+}
+
+export const StridedSlice = 'StridedSlice';
+export type StridedSliceInputs = Pick<NamedTensorInfoMap, 'x'>;
+export interface StridedSliceAttrs {
+  begin: number[];
+  end: number[];
+  strides: number[];
+  beginMask: number;
+  endMask: number;
+  ellipsisMask: number;
+  newAxisMask: number;
+  shrinkAxisMask: number;
+}
+
 export const Tan = 'Tan';
 export type TanInputs = UnaryInputs;
 
@@ -658,6 +690,13 @@ export interface TileAttrs {
   reps: number[];
 }
 
+export const TopK = 'TopK';
+export type TopKInputs = Pick<NamedTensorInfoMap, 'x'>;
+export interface TopKAttrs {
+  k: number;
+  sorted: boolean;
+}
+
 export const Transpose = 'Transpose';
 export type TransposeInputs = Pick<NamedTensorInfoMap, 'x'>;
 export interface TransposeAttrs {

tfjs-core/src/ops/log_softmax.ts

@@ -0,0 +1,80 @@
+/**
+ * @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 {ENGINE, ForwardFunc} from '../engine';
+import {LogSoftmax, LogSoftmaxAttrs, LogSoftmaxInputs} 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 {max} from './max';
+import {op} from './operation';
+import {sub} from './sub';
+
+/**
+ * Computes the log softmax.
+ *
+ * ```js
+ * const a = tf.tensor1d([1, 2, 3]);
+ *
+ * a.logSoftmax().print();  // or tf.logSoftmax(a)
+ * ```
+ *
+ * ```js
+ * const a = tf.tensor2d([2, 4, 6, 1, 2, 3], [2, 3]);
+ *
+ * a.logSoftmax().print();  // or tf.logSoftmax(a)
+ * ```
+ *
+ * @param logits The logits array.
+ * @param axis The dimension softmax would be performed on. Defaults to `-1`
+ *     which indicates the last dimension.
+ */
+/** @doc {heading: 'Operations', subheading: 'Normalization'} */
+function logSoftmax_<T extends Tensor>(logits: T|TensorLike, axis = -1): T {
+  const $logits = convertToTensor(logits, 'logits', 'logSoftmax');
+
+  if (axis === -1) {
+    axis = $logits.rank - 1;
+  }
+  if (axis !== $logits.rank - 1) {
+    throw Error(
+        'Log Softmax along a non-last dimension is not yet supported. ' +
+        `Logits was rank ${$logits.rank} and axis was ${axis}`);
+  }
+
+  const forward: ForwardFunc<Tensor> = (backend, save) => {
+    const keepDims = true;
+    const xMax = max(logits, axis, true);
+    const shifted = sub(logits, xMax);
+    const value =
+        shifted.toFloat().sub(shifted.exp().sum(axis, keepDims).log());
+    save([value]);
+    return value;
+  };
+
+  const inputs: LogSoftmaxInputs = {logits: $logits};
+  const attrs: LogSoftmaxAttrs = {axis};
+
+  return ENGINE.runKernelFunc(
+             forward, inputs as {} as NamedTensorMap, null /* grad */,
+             LogSoftmax, attrs as {} as NamedAttrMap) as T;
+}
+
+export const logSoftmax = op({logSoftmax_});

tfjs-core/src/ops/log_softmax_test.ts

@@ -0,0 +1,84 @@
+/**
+ * @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 * as tf from '../index';
+import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
+import {expectArraysClose} from '../test_util';
+
+describeWithFlags('logSoftmax', ALL_ENVS, () => {
+  it('regular test', async () => {
+    const y = tf.logSoftmax(tf.tensor1d([2, 1, 3]));
+
+    expectArraysClose(await y.data(), [-1.407606, -2.4076061, -0.407606]);
+  });
+
+  it('Huge difference', async () => {
+    const y = tf.logSoftmax(tf.tensor1d([-1000, +1000]));
+
+    expectArraysClose(await y.data(), [-2000, 0]);
+  });
+
+  it('Propagates NaNs', async () => {
+    const a = tf.tensor1d([2, 1, NaN]);
+    const y = tf.logSoftmax(a);
+    expectArraysClose(await y.data(), [NaN, NaN, NaN]);
+  });
+
+  it('2D, axis=1', async () => {
+    const y = tf.logSoftmax(tf.tensor2d([[2, 1, 3], [1, 3, 2]], [2, 3]), 1);
+    const expected =
+        [-1.407606, -2.4076061, -0.407606, -2.4076061, -0.4076061, -1.4076061];
+    expect(y.rank).toBe(2);
+    expectArraysClose(await y.data(), expected);
+  });
+
+  it('2D, implicit axis=1', async () => {
+    const y = tf.logSoftmax(tf.tensor2d([[2, 1, 3], [1, 3, 2]], [2, 3]));
+    const expected =
+        [-1.407606, -2.4076061, -0.407606, -2.4076061, -0.4076061, -1.4076061];
+    expect(y.rank).toBe(2);
+    expectArraysClose(await y.data(), expected);
+  });
+
+  it('1D gradient', async () => {
+    const x = tf.tensor1d([1, 2, 10]);
+    const dy = tf.tensor1d([1, 2, 3]);
+    const dx = tf.grad((x) => x.logSoftmax())(x, dy);
+
+    expect(dx.shape).toEqual(x.shape);
+    expectArraysClose(await dx.data(), [0.9992599, 1.9979881, -2.9972477]);
+  });
+
+  it('2D, axis=0 throws error', () => {
+    const f = () => {
+      tf.logSoftmax(tf.tensor2d([[2, 1, 3], [1, 3, 2]], [2, 3]), 0);
+    };
+    expect(f).toThrowError();
+  });
+
+  it('throws when passed a non-tensor', () => {
+    expect(() => tf.logSoftmax({} as tf.Tensor))
+        .toThrowError(
+            /Argument 'logits' passed to 'logSoftmax' must be a Tensor/);
+  });
+
+  it('accepts a tensor-like object', async () => {
+    const y = tf.logSoftmax([2, 1, 3]);
+
+    expectArraysClose(await y.data(), [-1.407606, -2.4076061, -0.407606]);
+  });
+});

tfjs-core/src/ops/ops.ts

@@ -87,6 +87,7 @@ export {localResponseNormalization} from './local_response_normalization';
 export {log} from './log';
 export {log1p} from './log1p';
 export {logSigmoid} from './log_sigmoid';
+export {logSoftmax} from './log_softmax';
 export {logSumExp} from './log_sum_exp';
 export {logicalAnd} from './logical_and';
 export {logicalNot} from './logical_not';
@@ -149,12 +150,14 @@ export {slice1d} from './slice1d';
 export {slice2d} from './slice2d';
 export {slice3d} from './slice3d';
 export {slice4d} from './slice4d';
+export {softmax} from './softmax';
 export {spaceToBatchND} from './space_to_batch_nd';
 export {split} from './split';
 export {square} from './square';
 export {squaredDifference} from './squared_difference';
 export {squeeze} from './squeeze';
 export {stack} from './stack';
+export {stridedSlice} from './strided_slice';
 export {sub} from './sub';
 export {sum} from './sum';
 export {tan} from './tan';
@@ -167,6 +170,7 @@ export {tensor4d} from './tensor4d';
 export {tensor5d} from './tensor5d';
 export {tensor6d} from './tensor6d';
 export {tile} from './tile';
+export {topk} from './topk';
 export {truncatedNormal} from './truncated_normal';
 export {unsortedSegmentSum} from './unsorted_segment_sum';
 export {unstack} from './unstack';
@@ -181,11 +185,8 @@ export * from './unary_ops';
 export * from './compare';
 export * from './binary_ops';
 export * from './transpose';
-export * from './softmax';
 export * from './norm';
 export * from './moving_average';
-export * from './strided_slice';
-export * from './topk';
 export * from './scatter_nd';
 export * from './spectral_ops';
 export * from './sparse_to_dense';