Modularize cast, complex, concat, add, sub, multiply, notEqual, slice, fft, ifft.

[lina128]

This PR has below changes:

1. It constitutes a minimal set of kernels need to modularize together due to Complex kernel modularization.
2. Changed Complex underlying tensors to tensorInfos.
3. Created/Refactored utility functions for binary kernels.
4. Moved complex related memory usage test to cpu and webgl backend, as they are backend specific now.

---

This change is 

[tafsiri]

Great work on this PR. There are probably a few things that might warrant offline discussion but take a look through first and we can discuss with @annxingyuan here first and then chat if needed.

Reviewed 33 of 33 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

tfjs-backend-cpu/src/backend_cpu.ts, line 56 at r1 (raw file):

  // individual tensors, with a parent joining the two with the
  // complexTensors field.
  complexTensors?: {real: TensorInfo, imag: TensorInfo};

should we rename this to complexTensorData or complexTensorInfo

---

tfjs-backend-cpu/src/backend_cpu.ts, line 139 at r1 (raw file):

      const realValues =
          this.readSync(complexTensors.real.dataId) as Float32Array;

nit: remove this newline and the one right below.

---

tfjs-backend-cpu/src/backend_cpu.ts, line 1211 at r1 (raw file):

  }

  int<T extends Tensor>(x: T): T {

this may still need to be a method on backend (it is like a utility/internal implementation detail—that seems to be used only by cast?). I see in this PR you have converted it to a kernel. Is there a corresponding tensorflow kernel?

---

tfjs-backend-cpu/src/kernels/Add.ts, line 20 at r1 (raw file):

import {Add, KernelConfig} from '@tensorflow/tfjs-core';

import {createBinaryKernelComplexSupportImpl, createBinaryKernelComplexSupportKernelFunc, createBinaryKernelImpl} from '../utils/kernel_utils';

naming suggestion for these utilities (that basically mirrors the old names)

createBinaryKernelImpl -> broadcastedBinaryKernelSimple
createBinaryKernelComplexSupportImpl = broadcastedBinaryKernelComplex
createBinaryKernelComplexSupportKernelFunc -> binaryKernelFunc (and make complexImpl an optional parameter). This will allow you to do warnings about whether a kernel supports complex params in one place.

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 34 at r1 (raw file):

  const {dtype} = attrs;

  if (dtype === 'complex64') {

Suggested comment line above this block. // casting to complex64

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 40 at r1 (raw file):

    // TODO(lina128): Import kernel function once zeros is modularized.
    const zerosTensor = tf.zeros(x.shape);

you could use makeTensorInfo here right? Would the downside of that be that it is effectively the implementation of zeros? on the other hand it could pretty much stay that way (it doesn't have to import the zeros kernel func). I'm not necessarily opposed to this approach (and then replacing it with the kernelFunc call) but wanted to bring it up for discussion.

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 58 at r1 (raw file):

    return {dataId: result.dataId, shape: result.shape, dtype};
  }

Suggested comment line above this block. // casting from complex64

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 59 at r1 (raw file):

  }

  if (x.dtype === 'complex64') {

Can we ever cast from complex64 to anything else without encoding loss? I ask this because I wanted to suggest moving this up to come right after the block on after the block on casting to complex64. Reading hasEncodingLoss suggests that if you cast from complex64 to bool you wouldn't get encoding loss, but i don't think that is true. Am I missing something?

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 69 at r1 (raw file):

  if (dtype === 'int32') {
    return int({inputs: {x}, backend});

this is what i suspect should be backend.int() or just makeTensorInfo with the existing vals and an int dtype.

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 76 at r1 (raw file):

    // modularized.
    const zero = tf.scalar(0, x.dtype);

nit: remove this newline

---

tfjs-backend-cpu/src/kernels/Complex.ts, line 34 at r1 (raw file):

  const complex = backend.data.get(complexInfo.dataId);

  // The complex tensor owns the underlying real and imag tensors, only the

tensors => tensorInfos

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 65 at r1 (raw file):

    expect(result.shape).toEqual([1]);
    expectArraysClose(await result.data(), [4, 6]);

could you add an assertion for expected numDataBuffers at this point.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 90 at r1 (raw file):

    real.dispose();
    imag.dispose();
    expect(tf.memory().numTensors).toBe(startTensors);

could you add an assertion for expected numDataBuffers at this point.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 93 at r1 (raw file):

  });

  it('tf.complex disposing underlying tensors', async () => {

is this covered by a core test? If not i think this could be in core since it does not look at anything backend specific.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 108 at r1 (raw file):

    imag.dispose();

    // real and imag still exist.

is this comment correct? should it be complex still exists?

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 141 at r1 (raw file):

    expect(b.shape).toEqual([6]);
    expectArraysClose(await a.data(), await b.data());

could you split what you have below into another test (something like "dispose after reshape"), it would repeat the same operations as above but without the previous 'expects'

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 175 at r1 (raw file):

    expect(b.dtype).toBe('complex64');
    expectArraysClose(await a.data(), await b.data());

similar to above split this into two.

---

tfjs-backend-cpu/src/kernels/Int.ts, line 17 at r1 (raw file):

 * =============================================================================
 */

See comments above about whether we need this kernel.

---

tfjs-backend-cpu/src/kernels/Reshape.ts, line 30 at r1 (raw file):

  const {shape} = attrs;

  const xSize = util.sizeFromShape(x.shape);

While i don't think these assertions are terribly expensive, we may want to actually measure the impact in a real model since we generally want reshape to be as cheap as possible. This could be a todo to look into with assertions on and off. We could also put this one behind a debug flag if it is costly.

---

tfjs-backend-cpu/src/kernels/Slice.ts, line 39 at r1 (raw file):

  const isContinous = slice_util.isSliceContinous(x.shape, $begin, $size);

nit: seems like a lot of extra newlines in this block and the next. :)

---

tfjs-backend-cpu/src/utils/fft_utils.ts, line 290 at r1 (raw file):

  const $imagVals = cpuBackend.data.get($imag.dataId).values as Float32Array;

  cpuBackend.disposeIntermediateTensorInfo(evenRealInfo);

optional: You could make disposeIntermediateTensorInfo take an array of TensorInfo as well.

---

tfjs-backend-cpu/src/utils/kernel_utils.ts, line 18 at r1 (raw file):

 */

import {backend_util, BinaryInputs, DataType, KernelConfig, KernelFunc, NumericDataType, TypedArray, util} from '@tensorflow/tfjs-core';

meta comment about the naming of this file, should this be renamed to binary_kernel_utils.ts to discourage this from getting infinitely long as it gains support functions for various kinds of kernels?

---

tfjs-backend-cpu/src/utils/kernel_utils.ts, line 25 at r1 (raw file):

import {complex} from '../kernels/Complex';

export type SimpleBinaryOp = (a: number, b: number) => number;

These types (their names) confuse me somewhat. Maybe if this were called BinaryOperation/BinaryOperationComplex and the one below SimpleBinaryKernel or something it would help with the the overloading issue between mathematical operations and our ops/kernels.

I also had some naming suggestions for the utility functions below.

---

tfjs-backend-wasm/src/kernels/Reshape.ts, line 31 at r1 (raw file):

  const {shape} = attrs as {} as ReshapeAttrs;

  const xSize = util.sizeFromShape(x.shape);

similar concern here about keeping wasm super cheap. Any particular reason or situation for adding this?

---

tfjs-core/src/engine_test.ts, line 548 at r1 (raw file):

      });

      it('can execute op with data from mixed backends', async () => {

did this move? curious about this deletion. seems like a test for core coordinating across multiple backends.

---

tfjs-core/src/kernel_names.ts, line 361 at r1 (raw file):

export type ImagInputs = Pick<NamedTensorInfoMap, 'input'>;

export const Int = 'Int';

See earlier comments about Int

---

tfjs-core/src/ops/reshape.ts, line 63 at r1 (raw file):

  const forward: ForwardFunc<
      Tensor<R>> = (backend: KernelBackend, save: GradSaveFunc) => {
    shape = util.inferFromImplicitShape(shape, $x.size) as ShapeMap[R];

I think these assertions are ones we want to keep at the op level.

---

tfjs-core/src/ops/slice.ts, line 72 at r1 (raw file):

  const forward: ForwardFunc<Tensor> = (backend, save) => {
    const [begin_, size_] = slice_util.parseSliceParams($x, begin, size);
    slice_util.assertParamsValid($x, begin_, size_);

same as above should be kept at op level as part of input validation.

---

tfjs-core/src/ops/slice_util.ts, line 18 at r1 (raw file):

 */

import {TensorInfo} from '../kernel_registry';

To the degree that these are op utils they should take tensors and not tensorInfos. Probably related to some of the input validation changes. so we should prob discuss offline.

---

tfjs-layers/src/layers/core_test.ts, line 118 at r1 (raw file):

  });

  describe('tensor with seed', () => {

is this related?

---

tfjs-layers/src/layers/normalization_test.ts, line 672 at r1 (raw file):

});

describeMathCPUAndGPU('LayerNormalization Layer: Symbolic', () => {

is this related?

[lina128]

Hi Yannick, thank you for the detailed review! Please see my reply below. Happy to discuss in next meeting :)

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

tfjs-backend-cpu/src/backend_cpu.ts, line 56 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

should we rename this to complexTensorData or complexTensorInfo

vote for complexTensorInfo. I just realized that this is not handled when moving data from backend to backend. I expect when moving data, we move the underlying complexTensor too, but looking at our code, it seems we don't handle it at all. Or am I missing something?

---

tfjs-backend-cpu/src/backend_cpu.ts, line 139 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

nit: remove this newline and the one right below.

Done.

---

tfjs-backend-cpu/src/backend_cpu.ts, line 1211 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

this may still need to be a method on backend (it is like a utility/internal implementation detail—that seems to be used only by cast?). I see in this PR you have converted it to a kernel. Is there a corresponding tensorflow kernel?

There's not a corresponding tensorflow kernel. And it is only used by cast. I can move it back to backend or move it to a util so that it is treeshakable.

---

tfjs-backend-cpu/src/kernels/Add.ts, line 20 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

naming suggestion for these utilities (that basically mirrors the old names)

createBinaryKernelImpl -> broadcastedBinaryKernelSimple
createBinaryKernelComplexSupportImpl = broadcastedBinaryKernelComplex
createBinaryKernelComplexSupportKernelFunc -> binaryKernelFunc (and make complexImpl an optional parameter). This will allow you to do warnings about whether a kernel supports complex params in one place.

Done.

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 34 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Suggested comment line above this block. // casting to complex64

Done.

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 58 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Suggested comment line above this block. // casting from complex64

Done.

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 59 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Can we ever cast from complex64 to anything else without encoding loss? I ask this because I wanted to suggest moving this up to come right after the block on after the block on casting to complex64. Reading hasEncodingLoss suggests that if you cast from complex64 to bool you wouldn't get encoding loss, but i don't think that is true. Am I missing something?

Good point. Moved up. I think it will always be encoding loss from complex64 to other types.

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 76 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

nit: remove this newline

Done.

---

tfjs-backend-cpu/src/kernels/Complex.ts, line 34 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

tensors => tensorInfos

Done.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 65 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

could you add an assertion for expected numDataBuffers at this point.

Done. Changed to test numDataIds.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 90 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

could you add an assertion for expected numDataBuffers at this point.

Done.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 93 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

is this covered by a core test? If not i think this could be in core since it does not look at anything backend specific.

Added numDataIds assertions.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 108 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

is this comment correct? should it be complex still exists?

Nice catch. Removed.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 141 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

could you split what you have below into another test (something like "dispose after reshape"), it would repeat the same operations as above but without the previous 'expects'

What's the reason? The test case follows the same pattern as other test in this file.

---

tfjs-backend-cpu/src/kernels/Reshape.ts, line 30 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

While i don't think these assertions are terribly expensive, we may want to actually measure the impact in a real model since we generally want reshape to be as cheap as possible. This could be a todo to look into with assertions on and off. We could also put this one behind a debug flag if it is costly.

Added TODO.

---

tfjs-backend-cpu/src/kernels/Slice.ts, line 39 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

nit: seems like a lot of extra newlines in this block and the next. :)

Removed newlines.

---

tfjs-backend-cpu/src/utils/kernel_utils.ts, line 18 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

meta comment about the naming of this file, should this be renamed to binary_kernel_utils.ts to discourage this from getting infinitely long as it gains support functions for various kinds of kernels?

I'm thinking to refactor this file in another PR to make this PR not getting even bigger. Agree that this should be binary_util.ts

---

tfjs-backend-cpu/src/utils/kernel_utils.ts, line 25 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

These types (their names) confuse me somewhat. Maybe if this were called BinaryOperation/BinaryOperationComplex and the one below SimpleBinaryKernel or something it would help with the the overloading issue between mathematical operations and our ops/kernels.

I also had some naming suggestions for the utility functions below.

Done.

---

tfjs-backend-wasm/src/kernels/Reshape.ts, line 31 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

similar concern here about keeping wasm super cheap. Any particular reason or situation for adding this?

We move the value transformation from reshape.ts op to kernels, these are moved along with them. So they are not getting more expensive than before. If we remove these checks, then there is no where this will be checked.

---

tfjs-core/src/engine_test.ts, line 548 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

did this move? curious about this deletion. seems like a test for core coordinating across multiple backends.

Ah my bad, I wanted to move this to e2e but forgot. Just moved to e2e backend_tests.ts. Keep the test here is kind of more complicated to set up now, because we also have to register kernels for these fake backends. The main thing to test here is backend switch but not registering kernels, so I just move to e2e, the backend_tests.ts is all about testing switching backends. WDYT?

---

tfjs-core/src/ops/reshape.ts, line 63 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I think these assertions are ones we want to keep at the op level.

These have transformation, not just assertion. If not implemented in kernel, other kernels calling into Reshape kernel will fail (e.g. when the shape has -1).

---

tfjs-core/src/ops/slice.ts, line 72 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

same as above should be kept at op level as part of input validation.

These are not validating raw input, these are validating computed value.

---

tfjs-layers/src/layers/core_test.ts, line 118 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

is this related?

Yes, after slice get modularized, this will alert kernels not registered without this change. After wrapping in it(){}, it starts working again. We probably should change it to describeMathCPUAndGPU to make the dependency clearer.

---

tfjs-layers/src/layers/normalization_test.ts, line 672 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

is this related?

Yes, the test below needs to wait for backend to be ready and kernels to be registered.

[annxingyuan]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

e2e/integration_tests/backends_test.ts, line 130 at r2 (raw file):

    tfc.setBackend('webgl');
    // This scalar lives in cpu2.

should this be "this scalar lives in webgl"?

---

e2e/integration_tests/backends_test.ts, line 143 at r2 (raw file):

    expect(tfc.memory().numTensors).toBe(numTensors + 2);
    expect(tfc.findBackend('webgl').numDataIds()).toBe(webglNumDataIds + 2);
    expect(tfc.findBackend('cpu').numDataIds()).toBe(cpuNumDataIds);

i actually thought that both backends would have one more dataId than before :) could you briefly explain the expectations for cpu / webgl?

---

tfjs-backend-cpu/src/backend_cpu.ts, line 98 at r2 (raw file):

  }

  makeTensorInfo(

would it make sense for the interface here to just be shape and dtype, to match webgl?

---

tfjs-backend-cpu/src/kernels/Complex.ts, line 38 at r2 (raw file):

  // underlying tensorData will be disposed.
  complex.complexTensors = {
    real: backend.makeTensorInfo(realVals, real.shape, 'float32'),

so with your change, rather than sharing a dataId from the input real/imag (cloning), we copy the input values into the new tensorinfos - is that right? if so, could you help me understand why it's necessary to create a copy rather than sharing? sorry if i'm missing something we've already talked about!

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 71 at r2 (raw file):

    // They should not increase tensor count, only complex tensor does.
    // 3 new tensorData is created for complex, real and imag.
    expect(tf.memory().numTensors).toBe(numTensors + 1);

for my own understanding, did this behavior change? without your changes, would numTensors have increased by 3 here?

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 193 at r2 (raw file):

    const b = a.clone();

    // 1 new tensor from the reshape.

clone

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 195 at r2 (raw file):

    // 1 new tensor from the reshape.
    expect(tf.memory().numTensors).toBe(memoryBefore.numTensors + 2);
    // No new tensor data.

this comment should be updated

---

tfjs-backend-cpu/src/kernels/Imag.ts, line 32 at r2 (raw file):

  // When complex tensor is disposed, its underlying parts will be disposed too.
  // Make new tensor out of the imag value of the complex. This makes sure the
  // value is still accessible even if complex tensor is disposed.

just to make sure i'm understanding your comment correctly - do we not get this behavior today with clone?

---

tfjs-backend-wasm/src/kernels/Reshape.ts, line 37 at r2 (raw file):

      xSize === util.sizeFromShape($shape),
      () => `new shape: ${$shape}, old shape: ${x.shape}. New shape and old ` +
          `shape must have the same number of elements.`);

The new shape (${$shape}}) has ${util.sizefromShape($shape)} elements and the old shape (${x.shape}) has ${xSize} elements. The new shape and old shape must have the same number of elements.

---

tfjs-backend-webgl/src/kernels/Complex_test.ts, line 24 at r2 (raw file):

// tslint:disable-next-line: no-imports-from-dist
import {describeWithFlags, ALL_ENVS} from '@tensorflow/tfjs-core/dist/jasmine_util';

nit: it would be nice if we could compare these tests line by line across webgl / cpu

---

tfjs-backend-webgl/src/kernels/Complex_test.ts, line 52 at r2 (raw file):

    expect(tf.memory().numDataBuffers).toBe(numBuffers + 3);
    numTensors = tf.memory().numTensors;
    numBuffers = tf.memory().numDataBuffers;

iiuc, it looks like a consequence of using tensorinfos instead of tensors for complex tensors on the cpu is that we lose the ability to share data buckets for complex's real / imag components, rather we need to copy them into new tensorinfos? accordingly, do you expect these assertions to change after complex is modularized in webgl?

[lina128]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

e2e/integration_tests/backends_test.ts, line 130 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

should this be "this scalar lives in webgl"?

Done.

---

e2e/integration_tests/backends_test.ts, line 143 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

i actually thought that both backends would have one more dataId than before :) could you briefly explain the expectations for cpu / webgl?

Added more assertions. The result is disposed after engine.endScope().

---

tfjs-backend-cpu/src/kernels/Complex.ts, line 38 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

so with your change, rather than sharing a dataId from the input real/imag (cloning), we copy the input values into the new tensorinfos - is that right? if so, could you help me understand why it's necessary to create a copy rather than sharing? sorry if i'm missing something we've already talked about!

Because we can no longer 'keep' the underlying tensor, which means if complex is disposed, the underlying parts gets disposed too, imagine these two scenarios: (1) real and imag input tensors are disposed (2) extract the real part from complex and dispose complex, in both scenarios, we want to keep the real part. So we need to copy the value into new tensorinfo when creating complex and also copy the value into new tensorinfo in real kernel.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 71 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

for my own understanding, did this behavior change? without your changes, would numTensors have increased by 3 here?

Yes. After the change, only the complex tensor is exposed, the underlying parts are tensorData wrapped as tensorInfos, they are not exposed or can be used immediately.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 193 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

clone

Done.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 195 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

this comment should be updated

Why? No new tensor data is created.

[lina128]

Hi @annxingyuan and @tafsiri , thank you for your detailed review, I did another pass. Please take a look. Thank you!

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

tfjs-backend-cpu/src/backend_cpu.ts, line 56 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

vote for complexTensorInfo. I just realized that this is not handled when moving data from backend to backend. I expect when moving data, we move the underlying complexTensor too, but looking at our code, it seems we don't handle it at all. Or am I missing something?

Changed to complexTensorInfo

---

tfjs-backend-cpu/src/backend_cpu.ts, line 1211 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

There's not a corresponding tensorflow kernel. And it is only used by cast. I can move it back to backend or move it to a util so that it is treeshakable.

Changed to a util.

---

tfjs-backend-cpu/src/backend_cpu.ts, line 98 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

would it make sense for the interface here to just be shape and dtype, to match webgl?

Hi Ann, I changed the function name to makeTensorInfoData to make it less ambiguous. Actually curious why webgl's makeTensorInfo is implemented this way, i.e. write to the backend without value and then read out the tensorData and then assign the value. This adds two additional line of code comparing to makeTensorInfoData here. Is there a case like it needs to know the id before computing the value?

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 40 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

you could use makeTensorInfo here right? Would the downside of that be that it is effectively the implementation of zeros? on the other hand it could pretty much stay that way (it doesn't have to import the zeros kernel func). I'm not necessarily opposed to this approach (and then replacing it with the kernelFunc call) but wanted to bring it up for discussion.

Whether we directly use the implementation of zeros here or call zeros kernel has the same effect. I want to use kernel where-ever possible, because I think this force us to think how a large computation can be represented by a set of smaller computations. I remember that Ping mentioned something like eventually any computation can be represented by a minimum set of kernels, so if there is a day when we want to represent any logic by that set, this helps that, but it's a long stretch :)

---

tfjs-backend-cpu/src/kernels/Cast.ts, line 69 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

this is what i suspect should be backend.int() or just makeTensorInfo with the existing vals and an int dtype.

Changed int to a cast util.

---

tfjs-backend-cpu/src/kernels/Imag.ts, line 32 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

just to make sure i'm understanding your comment correctly - do we not get this behavior today with clone?

Discussed offline, today we use clone and keep, in the new world, we copy the value and create new data bucket.

---

tfjs-backend-cpu/src/kernels/Int.ts, line 17 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

See comments above about whether we need this kernel.

Removed.

---

tfjs-backend-cpu/src/kernels/Reshape.ts, line 30 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Added TODO.

Discussed offline, transformation validation should happen in kernels.

---

tfjs-backend-cpu/src/utils/fft_utils.ts, line 290 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

optional: You could make disposeIntermediateTensorInfo take an array of TensorInfo as well.

Got it. I want to leave it as is for now, and then after modularization, we can decide what's the best way to address something like this. My thinking is, if it happens rather often, we may consider automatic disposal, if it only happens once or twice, then optimize it or not doesn't make much difference.

---

tfjs-backend-wasm/src/kernels/Reshape.ts, line 37 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

The new shape (${$shape}}) has ${util.sizefromShape($shape)} elements and the old shape (${x.shape}) has ${xSize} elements. The new shape and old shape must have the same number of elements.

Done.

---

tfjs-backend-webgl/src/kernels/Complex_test.ts, line 52 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

iiuc, it looks like a consequence of using tensorinfos instead of tensors for complex tensors on the cpu is that we lose the ability to share data buckets for complex's real / imag components, rather we need to copy them into new tensorinfos? accordingly, do you expect these assertions to change after complex is modularized in webgl?

Discussed offline, depending on webgl implementation, the test may be different. Having the same implementation as cpu is the easiest when come to disposing, but it does require the underlying data to be copied to new tensorinfos.

---

tfjs-core/src/kernel_names.ts, line 361 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

See earlier comments about Int

Removed.

---

tfjs-core/src/ops/slice_util.ts, line 18 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

To the degree that these are op utils they should take tensors and not tensorInfos. Probably related to some of the input validation changes. so we should prob discuss offline.

This is asserting transformed value, after today's discussion, we agreed that transformed value assertion should happen in kernel.

[tafsiri]

Thanks Na! Looking good, left a few comments, but mostly about naming :).

Reviewed 5 of 13 files at r2, 3 of 4 files at r3, 12 of 12 files at r4.
Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

tfjs-backend-cpu/src/backend_cpu.ts, line 98 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Hi Ann, I changed the function name to makeTensorInfoData to make it less ambiguous. Actually curious why webgl's makeTensorInfo is implemented this way, i.e. write to the backend without value and then read out the tensorData and then assign the value. This adds two additional line of code comparing to makeTensorInfoData here. Is there a case like it needs to know the id before computing the value?

to chime in i found the previous name less ambiguous :). The new name makes me wonder if there is makeTensorInfoWithout data. But i'll let ya'll discuss and see what convention makes most sense.

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 141 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

What's the reason? The test case follows the same pattern as other test in this file.

Mostly granularity of tests, having fewer expectations per test makes it easier to track down potential source of bugs imo. In this case when reading this test it seems like the first half asserts that reshape works, and the bottom parts asserts that disposing a reshaped complex tensor work. The second half seems to be more of a test of disposing rather than reshaping.

I made a similar comment about the clone test below. I think reading the name of the test makes me think it should just focus on that particular operation.

What do you think?

---

tfjs-backend-cpu/src/kernels/Imag.ts, line 32 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Discussed offline, today we use clone and keep, in the new world, we copy the value and create new data bucket.

Question/thought for future consideration (and to reify my understanding), once we have complex/real/imag implemented in all the backends, we might be able to special case complex tensor handling to not force delete underlying tensorinfos from core and thus enable memory sharing? (I just took a look at the node backend and it's going to be interesting when accounting for mem management there...). My understanding is that with the need to support both approaches (where core has more direct interactions with backend data via disposeData) we need to copy.

No need to respond to this unless my understanding is terribly wrong :), just wanted to document this as a something we might be able to revisit if we want to (though I also think it's an okay strategy for cpu backend as is).

---

tfjs-backend-cpu/src/utils/fft_utils.ts, line 290 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Got it. I want to leave it as is for now, and then after modularization, we can decide what's the best way to address something like this. My thinking is, if it happens rather often, we may consider automatic disposal, if it only happens once or twice, then optimize it or not doesn't make much difference.

+1. tbh i think its probably okay if kernels are a bit more verbose/low level. less abstraction to worry about when thinking about performance.

---

tfjs-backend-cpu/src/utils/kernel_utils.ts, line 18 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

I'm thinking to refactor this file in another PR to make this PR not getting even bigger. Agree that this should be binary_util.ts

+1 for future PR.

---

tfjs-backend-cpu/src/utils/kernel_utils.ts, line 25 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Done.

I was also trying to suggest

SimpleBinaryOp => BinaryOperation
SimpleBinaryOpImpl => SimpleBinaryKernelImpl

ComplexBinaryOp => ComplexBinaryOperation
ComplexBinaryOpImpl => ComplexBinaryKernelImpl

What do you think?

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 22 at r4 (raw file):

import {assertNotComplex} from '../cpu_util';

export function int(args: {inputs: {x: TensorInfo}, backend: MathBackendCPU}):

Simplify this signature to (x: TensorInfo, backend: MathBackendCPU).

I would also call this method something like makeInt32TensorInfo()

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 27 at r4 (raw file):

  const {x} = inputs;

  assertNotComplex(x, 'int');

will this be necessary? given the logic already in cast?

A higher level question, would any other kernel use this method directly rather than calling cast? If not do we need to split this out from the cast kernel since its not a particularly long function.

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 33 at r4 (raw file):

Quoted 4 lines of code…

  const resultValues = new Int32Array(util.sizeFromShape(x.shape));
  for (let i = 0; i < values.length; ++i) {
    resultValues[i] = values[i];
  }

Replace this with Int32Array.from(values)?

[annxingyuan]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

e2e/integration_tests/backends_test.ts, line 159 at r4 (raw file):

  // tslint:disable-next-line: ban
  xit('can move complex tensor from cpu to webgl.', async () => {

does moving between backends not work? will we be able to re-enable these tests in the future?

---

tfjs-backend-cpu/src/backend_cpu.ts, line 98 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

to chime in i found the previous name less ambiguous :). The new name makes me wonder if there is makeTensorInfoWithout data. But i'll let ya'll discuss and see what convention makes most sense.

I think it's because webgl kernels operate on outputs that have been registered with texData (the backend data registry), but when we create these outputs we don't have the values yet.

makeTensorInfoWithData works for me, or maybe you could do makeTensorInfo but the values could be the final optional argument?

---

tfjs-backend-cpu/src/kernels/Complex.ts, line 38 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Because we can no longer 'keep' the underlying tensor, which means if complex is disposed, the underlying parts gets disposed too, imagine these two scenarios: (1) real and imag input tensors are disposed (2) extract the real part from complex and dispose complex, in both scenarios, we want to keep the real part. So we need to copy the value into new tensorinfo when creating complex and also copy the value into new tensorinfo in real kernel.

makes sense thanks for the explanation!

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 195 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Why? No new tensor data is created.

doesn't the assertion below say that there are three new dataIds? maybe i'm not sure what tensor data refers to in your comment.

---

tfjs-backend-webgl/src/kernels/Complex_test.ts, line 52 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Discussed offline, depending on webgl implementation, the test may be different. Having the same implementation as cpu is the easiest when come to disposing, but it does require the underlying data to be copied to new tensorinfos.

makes sense!

[tafsiri]

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

e2e/integration_tests/backends_test.ts, line 159 at r4 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

does moving between backends not work? will we be able to re-enable these tests in the future?

I had suggested this be skipped for now in this PR. My understanding is that this would fail on master (but was discovered while making these changes for CPU). Is that correct Na?

But yeah the idea was that we would have a marker that we have a skipped test to work on fixing.

[lina128]

Thank you again Yannick and Ann! Addressed the comment.

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

e2e/integration_tests/backends_test.ts, line 159 at r4 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I had suggested this be skipped for now in this PR. My understanding is that this would fail on master (but was discovered while making these changes for CPU). Is that correct Na?

But yeah the idea was that we would have a marker that we have a skipped test to work on fixing.

Yes, this would fail on master. Right now, webgl doesn't allow moving complex64 data. And moving data to cpu logic doesn't work due to readSync merge values. As Yannick suggested, I added this pair of tests. After webgl modularization, we should expect this be solved.

---

tfjs-backend-cpu/src/backend_cpu.ts, line 98 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

I think it's because webgl kernels operate on outputs that have been registered with texData (the backend data registry), but when we create these outputs we don't have the values yet.

makeTensorInfoWithData works for me, or maybe you could do makeTensorInfo but the values could be the final optional argument?

Ok, I will change it back to makeTensorInfo and make value optional. Thank you both for suggestion!

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 141 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Mostly granularity of tests, having fewer expectations per test makes it easier to track down potential source of bugs imo. In this case when reading this test it seems like the first half asserts that reshape works, and the bottom parts asserts that disposing a reshaped complex tensor work. The second half seems to be more of a test of disposing rather than reshaping.

I made a similar comment about the clone test below. I think reading the name of the test makes me think it should just focus on that particular operation.

What do you think?

The pattern I see in these memory test all start from some state (num of tensors, num of dataIds), then do something and assert along the way (with a set of assertions), at the end end with the same state as before. We can go extreme to create a test per assertion, which conforms with test principle, but we ended up having several tests with pretty much the same body, so maintenance becomes an issue. For debugging, it shouldn't be an issue, the failure will tell which assertion fail. Based on these two points, I prefer to keep the test this way instead of splitting. WDYT?

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 195 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

doesn't the assertion below say that there are three new dataIds? maybe i'm not sure what tensor data refers to in your comment.

Ah, I see the confusion. The numDataIdsBefore is the initial num of dataIds. The last assertion (3 lines above) also asserts numDataIdsBefore + 3, so this assertion just shows there is no change since last assertion.

---

tfjs-backend-cpu/src/kernels/Imag.ts, line 32 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Question/thought for future consideration (and to reify my understanding), once we have complex/real/imag implemented in all the backends, we might be able to special case complex tensor handling to not force delete underlying tensorinfos from core and thus enable memory sharing? (I just took a look at the node backend and it's going to be interesting when accounting for mem management there...). My understanding is that with the need to support both approaches (where core has more direct interactions with backend data via disposeData) we need to copy.

No need to respond to this unless my understanding is terribly wrong :), just wanted to document this as a something we might be able to revisit if we want to (though I also think it's an okay strategy for cpu backend as is).

Sounds good, once we modularized all backends, we can revisit. Basically, we need a solution to fulfill these two scenarios: (1) real and imag input tensors are disposed (2) extract the real part from complex and dispose complex.

---

tfjs-backend-cpu/src/utils/kernel_utils.ts, line 25 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I was also trying to suggest

SimpleBinaryOp => BinaryOperation
SimpleBinaryOpImpl => SimpleBinaryKernelImpl

ComplexBinaryOp => ComplexBinaryOperation
ComplexBinaryOpImpl => ComplexBinaryKernelImpl

What do you think?

Done. Except the first one I changed to SimpleBinaryOperation, for consistency.

---

tfjs-backend-webgl/src/kernels/Complex_test.ts, line 24 at r2 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

nit: it would be nice if we could compare these tests line by line across webgl / cpu

Acknowledged.

---

tfjs-layers/src/layers/core_test.ts, line 118 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Yes, after slice get modularized, this will alert kernels not registered without this change. After wrapping in it(){}, it starts working again. We probably should change it to describeMathCPUAndGPU to make the dependency clearer.

This describe is already in describeMathCPUAndGPU, so no need to change that part. We still have to wrap the test in it(){}, otherwise it cannot find the kernel. Something I think relate to the order of executing side effect vs. test runner execution order.

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 22 at r4 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Simplify this signature to (x: TensorInfo, backend: MathBackendCPU).

I would also call this method something like makeInt32TensorInfo()

Done.

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 27 at r4 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

will this be necessary? given the logic already in cast?

A higher level question, would any other kernel use this method directly rather than calling cast? If not do we need to split this out from the cast kernel since its not a particularly long function.

No other places use it right now, so removed the assertion. If the use case expand in future, we can move out at that time.

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 33 at r4 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

  const resultValues = new Int32Array(util.sizeFromShape(x.shape));
  for (let i = 0; i < values.length; ++i) {
    resultValues[i] = values[i];
  }

Replace this with Int32Array.from(values)?

Done.

[tafsiri]

LGTM from my end! Two suggestions though for you to consider.

Reviewed 10 of 10 files at r5.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

tfjs-backend-cpu/src/kernels/Complex_test.ts, line 141 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

The pattern I see in these memory test all start from some state (num of tensors, num of dataIds), then do something and assert along the way (with a set of assertions), at the end end with the same state as before. We can go extreme to create a test per assertion, which conforms with test principle, but we ended up having several tests with pretty much the same body, so maintenance becomes an issue. For debugging, it shouldn't be an issue, the failure will tell which assertion fail. Based on these two points, I prefer to keep the test this way instead of splitting. WDYT?

Sure, it's definitely somewhat subjective what the ideal granularity is. So +1 to your preference.

---

tfjs-core/src/engine_test.ts, line 548 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

Ah my bad, I wanted to move this to e2e but forgot. Just moved to e2e backend_tests.ts. Keep the test here is kind of more complicated to set up now, because we also have to register kernels for these fake backends. The main thing to test here is backend switch but not registering kernels, so I just move to e2e, the backend_tests.ts is all about testing switching backends. WDYT?

I do think its good for a basic version of this test to live here, for reference here is a code snippet for how we re-register all kernels to a new backend in tfjs-react-native

// Register all the webgl kernels on the rn-webgl backend
    const kernels = tf.getKernelsForBackend('webgl');
    kernels.forEach(kernelConfig => {
      const newKernelConfig =
          Object.assign({}, kernelConfig, {backendName: 'rn-webgl'});
      tf.registerKernel(newKernelConfig);
    });

In this case you can copy from 'cpu' to 'cpu1' and 'cpu2'.

Using fakeBackends means there isn't a dependency on any webgl specific behaviour (or bugs), it just focuses on move semantics and keeps it close to CPU.

So I would strongly recommend it (and this can be in addition to moving data between specific real backends in the e2e tests).

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 27 at r4 (raw file):

Previously, lina128 (Na Li) wrote…

No other places use it right now, so removed the assertion. If the use case expand in future, we can move out at that time.

Sorry by "do we need to split this out" i meant to ask, why split this out from cast at this time since nothing else is using it. Non blocking (but recommended), but I think you could just include this function (without exporting it), or just inline the functionality in cast. Right now it opens up the potential for using it in other functions whereas you may want to force them to go through cast.

[tafsiri]

And by suggestions I mean its okay if you don't take them :).

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

[tafsiri]

Also, great work on this PR @lina128!!

[lina128]

Reviewable status: complete! 2 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

tfjs-core/src/engine_test.ts, line 548 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I do think its good for a basic version of this test to live here, for reference here is a code snippet for how we re-register all kernels to a new backend in tfjs-react-native

// Register all the webgl kernels on the rn-webgl backend
    const kernels = tf.getKernelsForBackend('webgl');
    kernels.forEach(kernelConfig => {
      const newKernelConfig =
          Object.assign({}, kernelConfig, {backendName: 'rn-webgl'});
      tf.registerKernel(newKernelConfig);
    });

In this case you can copy from 'cpu' to 'cpu1' and 'cpu2'.

Using fakeBackends means there isn't a dependency on any webgl specific behaviour (or bugs), it just focuses on move semantics and keeps it close to CPU.

So I would strongly recommend it (and this can be in addition to moving data between specific real backends in the e2e tests).

Make sense. Added the test back, only test result, not memory.

---

tfjs-backend-cpu/src/utils/cast_utils.ts, line 27 at r4 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Sorry by "do we need to split this out" i meant to ask, why split this out from cast at this time since nothing else is using it. Non blocking (but recommended), but I think you could just include this function (without exporting it), or just inline the functionality in cast. Right now it opens up the potential for using it in other functions whereas you may want to force them to go through cast.

Ah, got it. Removed this util, just inline the code.