Add a complex64 dtype.

[nsthorat]

Adds a complex64 dtype to the Tensor class.

Adds support for:

• tf.complex, tf.real, tf.imag
• tf.reshape, tf.clone, tf.cast
• tf.add, tf.sub, tf.mul

In the WebGL backend, we add a new field to TextureData which points a DataId to an optional array of complex components (real and imaginary). These are stored as completely separate Tensors, and any operations on complex numbers will delegate to operations on the underlying Tensors.

I added a complicated operation, tf.mul, which generates real and imaginary parts and has to blend the real and imaginary parts of the input. This will serve as a starting point for the FFT PR.

---

This change is 

[tafsiri]

Wow, good stuff! I took a close look and left some comments. The only thing i didn't closely examine is broadcastedBinaryComplexOp as i may need to brush up a bit on broadcasting a bit more first and wanted to submit the rest of the comments.

Reviewed 24 of 25 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @nsthorat, @dsmilkov, and @tafsiri)

---

src/tensor.ts, line 45 at r1 (raw file):

    if (values != null) {
      const n = values.length;
      const size = util.sizeFromShape(shape, dtype);

I think I get this now, but i wonder if calling this something like storageSize or numValues would be clearer than size as size seems to mostly be used for something that doesn't change based on dtype.

It had initially stood out to me that other calls to util.sizeFromShape were not passing dtype and it looked a bit confusing in comparing the calls and expected return. I know it works but since both util.getArrayFromDType and util.sizeFromShape can use dtype information it can be . a bit confusing on first read.

Would you consider having a separate util.getStorageSize(shape, dtype) method that calls util.sizeFromShape(shape) (and removing its optional dtype param) to make the intent of the call clearer?

---

src/tensor.ts, line 84 at r1 (raw file):

    }

    const index = this.locToIndex(locs);

I might be missing something but doesn't locToIndex or computeStrides need to dtype into account now to not have the indices step on each other for the real & imaginary parts?

---

src/tensor.ts, line 108 at r1 (raw file):

    }
    return this.dtype === 'complex64' ?
        [this.values[index], this.values[index + 1]] :

Similar to above I would have thought the indexing would need to take dtype into account before the final index, index+1 return. I might be misunderstanding about the data layout or something else, I've just been thinking about it with a rank 1 tensor with 2 complex numbers.

I tried making little failing test, hopefully I'm using .set correctly

const buff = buffer([2], 'complex64');
expect(buff.shape).toEqual([2]);
buff.set([1, 2], 0);
buff.set([3, 4], 1);
console.log('buff', buff.toTensor().dataSync(), buff.values);
expectArraysEqual(Array.from(buff.toTensor().dataSync()), [1, 2, 3, 4]);

---

src/tensor_test.ts, line 1345 at r1 (raw file):

    expect(() => Tensor.make([0], {values}))
        .toThrowError(
            'Based on the provided shape, [0], and dtype float32, the tensor ' +

If making dtype explicit in the message would making dtype explicit in the call help to make sure they never skew? (Not that I anticipate us changing the default dtype).

---

src/types.ts, line 88 at r1 (raw file):

}

enum UpcastComplex64AndMap {

Not specific to this enum, but it might be nice to add a comment about these Upcast enums and what they represent.

---

src/util.ts, line 126 at r1 (raw file):

function deepAssertShapeConsistency(
    val: number|boolean|RegularArray<number>|RegularArray<boolean>,
    shape: number[], indices: number[], dtype: DataType) {

The function doesn't seem to use dtype other than pass it recursively.

---

src/util.ts, line 149 at r1 (raw file):

}

export function sizeFromShape(

Commented on this above, but I find that having the optional dtype makes the intent less clear to me.

---

src/kernels/backend_cpu.ts, line 184 at r1 (raw file):

    const result = complex_util.mergeRealAndImagArrays(
        real.dataSync() as Float32Array, imag.dataSync() as Float32Array);
    return tensor(result, real.shape, 'complex64') as T;

You could assert that real and imag have the same shape at the start.

---

src/kernels/backend_cpu.ts, line 197 at r1 (raw file):

  }

  private assertNotComplex(tensor: Tensor|Tensor[], opName: string) {

Could this be outside of a backend and used by ops as well? I'm thinking of things like tf.ones and tf.fill

---

src/kernels/backend_cpu.ts, line 308 at r1 (raw file):

    if (a.dtype === 'complex64' || b.dtype === 'complex64') {
      return this.broadcastedBinaryComplexOp(
                 a.cast('complex64'), b.cast('complex64'),

is this primarily to cast one of the tensors that might not be complex?

---

src/kernels/backend_webgl.ts, line 211 at r1 (raw file):

    const {texture, texShape, usage, dtype, shape} = texData;

    if (dtype === 'complex64') {

It might be worth adding a comment about the approach of making the backing tensors but deferring their underlying write until some op is called on them (my assumption/understanding from the code).

---

src/kernels/backend_webgl.ts, line 536 at r1 (raw file):

  multiply(a: Tensor, b: Tensor): Tensor {
    if (a.dtype === 'complex64') {

should this also check dtype of b to be complex

---

src/kernels/backend_webgl.ts, line 541 at r1 (raw file):

      const realProgram = new BinaryOpComplexProgram(
          binaryop_complex_gpu.COMPLEX_MULTIPLY.real, a.shape, b.shape);

The snake case feel a bit odd. Why not import COMPLEX_MULTIPLY directly as a named import. Or alternatively

import * as BinaryComplexGPU from ...
import {BinaryOpComplexProgram} from ...

---

src/kernels/backend_webgl.ts, line 1289 at r1 (raw file):

      const reshapedData = this.texData.get(reshapedTensor.dataId);
      reshapedData.complexTensors = {
        real: ENV.engine.keep(xData.complexTensors.real.reshape(shape)),

does the old (pre-reshape) real and imag tensors need to be disposed?

---

src/kernels/complex_util.ts, line 18 at r1 (raw file):

 */

export function mergeRealAndImagArrays(

Consider adding jsdoc

---

src/kernels/complex_util.ts, line 33 at r1 (raw file):

}

export function splitRealAndImagArrays(complex: Float32Array):

consider adding jsdoc

---

src/kernels/webgl/gpgpu_math.ts, line 43 at r1 (raw file):

}

export interface TensorData {

I'm a bit curious to know why the type param was removed. I'm not advocating either way nor sure what advantage it provided before.

---

src/kernels/webgl/tex_util.ts, line 39 at r1 (raw file):

  // individual tensors, with a parent joining the two with the
  // complexDataIds field. When this is defined, texture will be null.
  complexTensors: {real: Tensor, imag: Tensor};

should this explicitly be optional? or does that make the type checking annoying (checking for complexTensors vs dtype==='complex64')

---

src/ops/array_ops.ts, line 545 at r1 (raw file):

 * Given a tensor input, this operation returns a tensor of type float that is
 * the real part of each element in input considered as a complex number.
 *

Maybe mention what it does if input is real?

---

src/ops/tensor_ops.ts, line 361 at r1 (raw file):

/** @doc {heading: 'Tensors', subheading: 'Creation'} */
function ones<R extends Rank>(
    shape: ShapeMap[R], dtype: DataType = 'float32'): Tensor<R> {

I'm wondering if we need to limit this function (and others) to non complex data types. I can see why zeros can handle complex64 in a well defined way but I wonder if there may be some unexpected behaviour from ones and fill and others. On the other hand I think there would be an exception on tensor creation which would catch the problem, so it might not be an issue, but i thought i'd mention it as something to think about.

[dsmilkov]

Reviewed 16 of 25 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @nsthorat and @dsmilkov)

---

src/engine.ts, line 200 at r1 (raw file):

      this.numDataBuffers++;
      this.numBytes +=
          util.sizeFromShape(a.shape, a.dtype) * util.bytesPerElement(a.dtype);

Would there be problems with double counting bytes, since every complex tensor consists of real and img tensor whose bytes are already accounted for when the first ref count was made to them separately? Good to add a unit test for this.

---

src/tensor.ts, line 45 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I think I get this now, but i wonder if calling this something like storageSize or numValues would be clearer than size as size seems to mostly be used for something that doesn't change based on dtype.

It had initially stood out to me that other calls to util.sizeFromShape were not passing dtype and it looked a bit confusing in comparing the calls and expected return. I know it works but since both util.getArrayFromDType and util.sizeFromShape can use dtype information it can be . a bit confusing on first read.

Would you consider having a separate util.getStorageSize(shape, dtype) method that calls util.sizeFromShape(shape) (and removing its optional dtype param) to make the intent of the call clearer?

+1 for getStorageSize() to be separate from sizeFromShape().

---

src/tensor.ts, line 84 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I might be missing something but doesn't locToIndex or computeStrides need to dtype into account now to not have the indices step on each other for the real & imaginary parts?

+1. Unit tests for TensorBuffer with complex dtype would reveal this.

---

src/tensor.ts, line 400 at r1 (raw file):

      shape: ShapeMap[R], dtype: DataType, values?: TypedArray,
      dataId?: DataId) {
    this.size = util.sizeFromShape(shape);

The size property of tensor should match whatever tf.size(input) in TF Python returns, so we should make sure to align with TF Python for complex tensors.

---

src/tensor_test.ts, line 716 at r1 (raw file):

    expect(b.dtype).toBe('float32');
    expect(b.shape).toEqual([1, 1]);
    expectArraysClose(a.dataSync(), b.dataSync());

here and elsewhere, no need to call dataSync(). expectArrayClose should work with tensors

---

src/tensor_test.ts, line 734 at r1 (raw file):

    expect(b.shape).toEqual([1, 1]);
    expectArraysClose(a.dataSync(), b.dataSync());
  });

add a memory test with reshape, to make sure that there are no extra increfs happening, thus no memory leaks.

---

src/tensor_test.ts, line 914 at r1 (raw file):

    expect(result.dtype).toEqual('int32');
    expectArraysClose(result, [1, 2]);

Use expectArraysEqual() for int32

---

src/tensor_test.ts, line 922 at r1 (raw file):

    expect(result.dtype).toEqual('bool');
    expectArraysClose(result, [true, false]);

expectArraysEqual for 'bool'

---

src/tensor_test.ts, line 1395 at r1 (raw file):

            'Based on the provided shape, [1,3,0,5], and dtype float32, the ' +
            'tensor should have 0 values but has 3');
  });

Add more tests for complex dtype:

• squeeze
• toString()
• tensor with 0 in the shape

---

src/util.ts, line 96 at r1 (raw file):

export function inferShape(
    val: TypedArray|number|boolean|RegularArray<number>|RegularArray<boolean>,
    dtype: DataType): number[] {

I would keep inferShape agnostic of the dtype because when the user makes a tf.complex(real, imag) tensor, you can call inferShape separately on real and img.

---

src/util.ts, line 384 at r1 (raw file):

  }
  if (newType === 'float32' && oldType !== 'complex64') {
    return false;

add unit tests

---

src/kernels/backend_cpu.ts, line 184 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

You could assert that real and imag have the same shape at the start.

+1 but let's do that assertion in tf.complex(), instead of the backend for consistency - backends assume input has been sanitized - avoids code duplication across the 3 backends: webgl, cpu and TF C in node.

---

src/kernels/backend_cpu.ts, line 197 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Could this be outside of a backend and used by ops as well? I'm thinking of things like tf.ones and tf.fill

Input sanitation should live in a level higher, but in this case what this conveys is that for this particular CPU backend, bunch of ops don't support complex arithmetic yet.

---

src/kernels/backend_webgl.ts, line 43 at r1 (raw file):

import {AvgPool2DBackpropProgram} from './webgl/avg_pool_backprop_gpu';
import {BatchNormProgram} from './webgl/batchnorm_gpu';
import * as binaryop_complex_gpu from './webgl/binaryop_complex_gpu';

import the methods you need directly (at some point we should convert all * imports to do this)

---

src/kernels/backend_webgl.ts, line 211 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

It might be worth adding a comment about the approach of making the backing tensors but deferring their underlying write until some op is called on them (my assumption/understanding from the code).

The write is not deferred since Tensor.make(shape, {values: real}) will cause backend_webgl.write() to be called immediately, but +1 to leaving a one-liner comment what this does.

---

src/kernels/backend_webgl.ts, line 489 at r1 (raw file):

  }

  slice<T extends Tensor>(x: T, begin: number[], size: number[]): T {

add assertNotComplex() to a bunch of ops like you do in CPU backend.

---

src/kernels/backend_webgl.ts, line 536 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

should this also check dtype of b to be complex

I think a and b are guaranteed to be of the same dtype by higher level sanity checks.

---

src/kernels/backend_webgl.ts, line 541 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

The snake case feel a bit odd. Why not import COMPLEX_MULTIPLY directly as a named import. Or alternatively

import * as BinaryComplexGPU from ...
import {BinaryOpComplexProgram} from ...

+1

---

src/kernels/backend_webgl.ts, line 1282 at r1 (raw file):

  reshape<R extends Rank>(x: Tensor, shape: ShapeMap[R]): Tensor<R> {
    const reshapedTensor = backend_util.reshapeTensor(x, shape);

backend_util.reshapeTensor is a one-liner (had to remind myself what it does). More readable if you just call Tensor.make() directly here.

---

src/kernels/backend_webgl.ts, line 1289 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

does the old (pre-reshape) real and imag tensors need to be disposed?

I'm also a bit confused here regarding ref counting. Assume x ref count is 1 before this method gets called. Then the first line in this method will increase the ref count of x.dataId to 2 - ignore the dataIdof its real and img parts for a second. Am I understanding this correctly? So then, focusing on the real part, in order for the real part to get actually disposed, the ref count of x.real.dataId has to get to 0, which happens when the ref count of x.dataId gets to 0, right?

---

src/ops/arithmetic_test.ts, line 513 at r1 (raw file):

    expect(result.dtype).toBe('complex64');
    expect(result.shape).toEqual([1]);
    expectArraysClose(result, [2 * 4 - 3 * 5, 2 * 5 + 3 * 4]);

test also that there are no mem-leaks using tf.memory()

---

src/ops/arithmetic_test.ts, line 1014 at r1 (raw file):

    expect(result.dtype).toBe('complex64');
    expect(result.shape).toEqual([1]);
    expectArraysClose(result, [4, 6]);

check for mem leaks using tf.memory()

---

src/ops/arithmetic_test.ts, line 1344 at r1 (raw file):

  });

  it('complex number broadcasting addition', () => {

s/addition/subtraction/

---

src/ops/arithmetic_test.ts, line 1374 at r1 (raw file):

        .toThrowError(
            // tslint:disable-next-line:max-line-length
            /The dtypes of the first\(int32\) and second\(float32\) input must match/);

add one more incompatible dtypes check where one is complex

---

src/ops/array_ops.ts, line 516 at r1 (raw file):

 * Converts two real numbers to a complex number.
 *
 * Given a tensor real representing the real part of a complex number, and a

wrap real and img in ``

---

src/ops/array_ops.ts, line 518 at r1 (raw file):

 * Given a tensor real representing the real part of a complex number, and a
 * tensor imag representing the imaginary part of a complex number, this
 * operation returns complex numbers elementwise of the form , where a

missing an ending: of the form....

---

src/ops/array_ops.ts, line 528 at r1 (raw file):

 * const complex = tf.complex(real, imag);
 *
 * complex.print();

need to change tensor_format.ts to handle complex numbers.

---

src/ops/array_ops_test.ts, line 3577 at r1 (raw file):

    expect(tf.memory().numTensors).toBe(startTensors);
  });
});

I left another comment, but will leave one here as well - would be good to add memory tests with reshape(), clone().

---

src/ops/tensor_ops.ts, line 61 at r1 (raw file):

        'array of numbers or booleans, or a TypedArray');
  }
  const inferredShape = inferShape(values, dtype);

how about asking the user to call tf.complex() for complex dtypes instead of tf.tensor() by throwing an error if user passes complex as dtype? This way we don't have to complicate inferShape with dtype info, and also we don't have to explain to the user how to pack (interleaved vs concatenatd) complex numbers when calling tf.tensor().

---

src/ops/tensor_ops.ts, line 361 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I'm wondering if we need to limit this function (and others) to non complex data types. I can see why zeros can handle complex64 in a well defined way but I wonder if there may be some unexpected behaviour from ones and fill and others. On the other hand I think there would be an exception on tensor creation which would catch the problem, so it might not be an issue, but i thought i'd mention it as something to think about.

We should align with tf.ones/tf.zeros in python, so if they allow complex as dtype, so should we.

[nsthorat]

Reviewable status: 0 of 1 approvals obtained (waiting on @tafsiri and @nsthorat)

---

src/engine.ts, line 200 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Would there be problems with double counting bytes, since every complex tensor consists of real and img tensor whose bytes are already accounted for when the first ref count was made to them separately? Good to add a unit test for this.

Good catch. Unfortunately I had to wire a "countBytes" through.

---

src/tensor.ts, line 45 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

+1 for getStorageSize() to be separate from sizeFromShape().

Done.

---

src/tensor.ts, line 84 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

+1. Unit tests for TensorBuffer with complex dtype would reveal this.

The shapes don't actually reflect that there are 2x as many values, we index into the shape as usual (this is how TensorFlow does it as well). I didn't want to touch this code since it's sort of non-trivial and has a lot of side effects. The layout of the underlying Float32Array is just 2x as large, and we compute everything as if it was the 1x (and then use a (2 * idx, 2 * idx + 1) indexing scheme into the underlying memory.

I did screw this up though, added a unit test.

---

src/tensor.ts, line 108 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Similar to above I would have thought the indexing would need to take dtype into account before the final index, index+1 return. I might be misunderstanding about the data layout or something else, I've just been thinking about it with a rank 1 tensor with 2 complex numbers.

I tried making little failing test, hopefully I'm using .set correctly

const buff = buffer([2], 'complex64');
expect(buff.shape).toEqual([2]);
buff.set([1, 2], 0);
buff.set([3, 4], 1);
console.log('buff', buff.toTensor().dataSync(), buff.values);
expectArraysEqual(Array.from(buff.toTensor().dataSync()), [1, 2, 3, 4]);

Yep, I messed this up. Needed to be 2 * x + 1. Added unit tests. Thanks!

---

src/tensor.ts, line 400 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

The size property of tensor should match whatever tf.size(input) in TF Python returns, so we should make sure to align with TF Python for complex tensors.

Yep, I made sure this lines up. Added storageSize, which is 2x size.

---

src/tensor_test.ts, line 716 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

here and elsewhere, no need to call dataSync(). expectArrayClose should work with tensors

Unfortunately you can't, if both are tensors we check the shapes match.

---

src/tensor_test.ts, line 734 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

add a memory test with reshape, to make sure that there are no extra increfs happening, thus no memory leaks.

Done.

---

src/tensor_test.ts, line 914 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Use expectArraysEqual() for int32

Done.

---

src/tensor_test.ts, line 922 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

expectArraysEqual for 'bool'

Done.

---

src/tensor_test.ts, line 1345 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

If making dtype explicit in the message would making dtype explicit in the call help to make sure they never skew? (Not that I anticipate us changing the default dtype).

Done.

---

src/tensor_test.ts, line 1395 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Add more tests for complex dtype:

• squeeze
• toString()
• tensor with 0 in the shape

Done.

---

src/types.ts, line 88 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Not specific to this enum, but it might be nice to add a comment about these Upcast enums and what they represent.

Done.

---

src/util.ts, line 96 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

I would keep inferShape agnostic of the dtype because when the user makes a tf.complex(real, imag) tensor, you can call inferShape separately on real and img.

If you call tf.tensor([1, 1], 'complex64') I want this to just work (to mirror python world), which is why we need dtype.

---

src/util.ts, line 126 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

The function doesn't seem to use dtype other than pass it recursively.

Done.

---

src/util.ts, line 149 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Commented on this above, but I find that having the optional dtype makes the intent less clear to me.

Done, removed, and added storage size.

---

src/util.ts, line 384 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

add unit tests

Done.

---

src/kernels/backend_cpu.ts, line 184 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

You could assert that real and imag have the same shape at the start.

Done. Did this in array_ops.

---

src/kernels/backend_cpu.ts, line 197 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Input sanitation should live in a level higher, but in this case what this conveys is that for this particular CPU backend, bunch of ops don't support complex arithmetic yet.

tf.ones doesn't take a tensor, nor does tf.fill

It's backend specific about what we support, TensorFlow C API may support complex numbers everywhere so I didn't want to do that at a higher level.

---

src/kernels/backend_cpu.ts, line 308 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

is this primarily to cast one of the tensors that might not be complex?

Yep, exactly. It's free if it's already complex.

---

src/kernels/backend_webgl.ts, line 43 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

import the methods you need directly (at some point we should convert all * imports to do this)

Will do this across the board in a separate PR since I will have to change some of the naming so it reads nicely.

---

src/kernels/backend_webgl.ts, line 211 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

The write is not deferred since Tensor.make(shape, {values: real}) will cause backend_webgl.write() to be called immediately, but +1 to leaving a one-liner comment what this does.

That's not what's happening, here we take the written values and split them into real and imaginary parts, writing them to new tensors in the backend. It happens immediately.

Added a comment.

---

src/kernels/backend_webgl.ts, line 489 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

add assertNotComplex() to a bunch of ops like you do in CPU backend.

Don't need to. This happens automatically in compileAndRun (since you can't run a program over a complex64 tensor directly).

---

src/kernels/backend_webgl.ts, line 536 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

should this also check dtype of b to be complex

multiply is strict with dtypes (they must match)

---

src/kernels/backend_webgl.ts, line 541 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

+1

This is the convention throughout the rest of the the file, binaryop_complex_gpu is the filename and COMPLEX_MULTIPLY is like a string constant, see other examples like this:
https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/backend_webgl.ts#L487

I see your point though the "real" and "imag" are strange, I made "real" and "imag" upper case, so it's more consistent.

Happy to make the file level import change across the board in a follow up PR.

---

src/kernels/backend_webgl.ts, line 1282 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

backend_util.reshapeTensor is a one-liner (had to remind myself what it does). More readable if you just call Tensor.make() directly here.

I changed the way reshape works (it doesnt reshape the underlying tensors, rather now we have a TensorHandle which has a shape and a dataId, so we can marry those two outside of a Tensor).

---

src/kernels/backend_webgl.ts, line 1289 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

I'm also a bit confused here regarding ref counting. Assume x ref count is 1 before this method gets called. Then the first line in this method will increase the ref count of x.dataId to 2 - ignore the dataIdof its real and img parts for a second. Am I understanding this correctly? So then, focusing on the real part, in order for the real part to get actually disposed, the ref count of x.real.dataId has to get to 0, which happens when the ref count of x.dataId gets to 0, right?

You are all right, this implementation is wrong. I updated the implementation so that reshape is as before not mutating the underlying tensors. When a complex operation is executed, we use the logical shape of the complex tensor with the data of the underlying parts. Added unit test for reshape => op with complex dtypes.

Great catch.

---

src/kernels/complex_util.ts, line 18 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Consider adding jsdoc

Done.

---

src/kernels/complex_util.ts, line 33 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

consider adding jsdoc

Done.

---

src/kernels/webgl/gpgpu_math.ts, line 43 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

I'm a bit curious to know why the type param was removed. I'm not advocating either way nor sure what advantage it provided before.

It was completely useless :) there is no place in the code below that ties two types together.

---

src/kernels/webgl/tex_util.ts, line 39 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

should this explicitly be optional? or does that make the type checking annoying (checking for complexTensors vs dtype==='complex64')

Made it optional. No good reason for this way.

---

src/ops/arithmetic_test.ts, line 513 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

test also that there are no mem-leaks using tf.memory()

I have a whole test just for that, see the bottom of array_ops_test.ts. Memory is different for webgl backend vs everything else, so I can't sprinkle them everywhere.

---

src/ops/arithmetic_test.ts, line 1014 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

check for mem leaks using tf.memory()

Whole test for that at bottom of array_ops_test.ts

---

src/ops/arithmetic_test.ts, line 1344 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

s/addition/subtraction/

Done.

---

src/ops/arithmetic_test.ts, line 1374 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

add one more incompatible dtypes check where one is complex

Done.

---

src/ops/array_ops.ts, line 516 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

wrap real and img in ``

Done.

---

src/ops/array_ops.ts, line 518 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

missing an ending: of the form....

Done.

---

src/ops/array_ops.ts, line 528 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

need to change tensor_format.ts to handle complex numbers.

Done.

---

src/ops/array_ops.ts, line 545 at r1 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Maybe mention what it does if input is real?

Done.

---

src/ops/array_ops_test.ts, line 3577 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

I left another comment, but will leave one here as well - would be good to add memory tests with reshape(), clone().

Done.

---

src/ops/tensor_ops.ts, line 61 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

how about asking the user to call tf.complex() for complex dtypes instead of tf.tensor() by throwing an error if user passes complex as dtype? This way we don't have to complicate inferShape with dtype info, and also we don't have to explain to the user how to pack (interleaved vs concatenatd) complex numbers when calling tf.tensor().

This isn't how TensorFlow works, though. I feel like it's pretty nice to just be able to create a complex64 directly. Let me know if you feel strongly and I will remove this code (which isn't really that complex)

---

src/ops/tensor_ops.ts, line 361 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

We should align with tf.ones/tf.zeros in python, so if they allow complex as dtype, so should we.

Yep, I made sure this lines up with python, which supports ones and zeros.

[nsthorat]

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

---

src/tensor.ts, line 45 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

Done.

No more need for storageSize.

---

src/engine.ts, line 200 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

Good catch. Unfortunately I had to wire a "countBytes" through.

Removed countBytes in favor of forcing tf.complex() instead of allowing complex dtypes for tensor constructors.

[dsmilkov]

An amazing work. Congrats! Left a dozen comments, but nothing major at this point - this PR is almost there. LGTM

Reviewed 1 of 25 files at r1, 20 of 23 files at r2.
Reviewable status: 0 of 1 approvals obtained (waiting on @tafsiri, @nsthorat, and @dsmilkov)

---

src/tensor.ts, line 38 at r2 (raw file):

export class TensorBuffer<R extends Rank> {
  size: number;
  storageSize: number;

remove this property

---

src/tensor.ts, line 82 at r2 (raw file):

    const index = this.locToIndex(locs);
    this.values[index] = value as number;

no need to say as number

---

src/tensor.ts, line 389 at r2 (raw file):

   * numbers.
   */
  readonly storageSize: number;

remove

---

src/tensor.ts, line 412 at r2 (raw file):

    if (values != null) {
      util.assert(
          this.storageSize === values.length,

undo this assert now that storageSize is gone.

---

src/tensor.ts, line 521 at r2 (raw file):

   * @param locs The location indices.
   */
  get(...locs: number[]) {

this is getting into tricky API territory, but right now it's strange that get() will return a single number instead of a tuple and the indexing is implicitly halved. File an issue to track this , but don't handle it here (it's a rare use-case to use .get() on complex)

---

src/tensor_test.ts, line 1081 at r2 (raw file):

    // Bytes should be counted once.
    expect(tf.memory().numBytes)
        .toBe(memoryBefore.numBytes + 12 * BYTES_PER_COMPLEX_ELEMENT);

shouldn't this be 6 * BYTES_PER_COMPLEX, since at the end of the day, there are 12 floats (124 bytes), or 6 "complex floats" (68 bytes). Or I guess, what' happening is that numBytes doubles when registerTensor() gets called with the complex dtype. This is a failure of the current design and assumption of engine. Add a comment here why the number of bytes doubled. Useful when we revisit this test.

---

src/util.ts, line 161 at r2 (raw file):

}

export function getStorageSize(shape: number[], dtype: DataType) {

remove

---

src/util.ts, line 338 at r2 (raw file):

export function getTypedArrayFromDType<D extends DataType>(
    dtype: D, storageSize: number): DataTypeMap[D] {

undo the changes to this method. IIUC, you don't call this method with complex64 anymore.

---

src/kernels/backend_cpu.ts, line 184 at r2 (raw file):

    const values = complex_util.mergeRealAndImagArrays(
        real.dataSync() as Float32Array, imag.dataSync() as Float32Array);
    return Tensor.make(real.shape, {values}, 'complex64') as T;

backend_webl calls Tensor.make(real.shape, {}) and uses extra book-keeping to keep pointers to the real and img tensors. backend_cpu doesn't do this. Should we align them?

---

src/kernels/webgl/tex_util.ts, line 38 at r2 (raw file):

  // For complex numbers, the real and imaginary parts are stored as their own
  // individual tensors, with a parent joining the two with the
  // complexDataIds field. When this is defined, texture will be null.

complexDataIds field doesn't exist. did you mean complexTensors fields?

---

src/ops/arithmetic_test.ts, line 513 at r1 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

I have a whole test just for that, see the bottom of array_ops_test.ts. Memory is different for webgl backend vs everything else, so I can't sprinkle them everywhere.

Do you know why memory is different across cpu and webgl? Can we make those tests run on both cpu and webgl, but not node (easy to do by passing {IS_NODE: false} as test constraint).

---

src/ops/array_ops_test.ts, line 3570 at r2 (raw file):

});

describeWithFlags('complex64 memory webgl', WEBGL_ENVS, () => {

I feel like array_ops_test is not the right file for these tests. How about moving them to complex64_test.ts ?

---

src/ops/array_ops_test.ts, line 3633 at r2 (raw file):

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

Add one more unit test where you make a real, img, and complex, and dispose the real and img, without disposing the complex, and making sure dataSync() still works on the complex. Smth like:

real = tensor(...)
img = tensor(...)
complex = tf.complex(real, img);
real.dispose()
img.dispose();
// Test memory
// Test complex.dataSync()
complex.dispose();
// Test memory

[tafsiri]

Looks good, left a few small comments

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

---

src/tensor.ts, line 55 at r2 (raw file):

    }
    if (dtype === 'complex64') {
      throw new Error(

Nice simplification

---

src/kernels/backend_cpu.ts, line 2335 at r2 (raw file):

      beta: number): Tensor4D {
    this.assertNotComplex(dy, 'LRNGrad');
    const channels = dy.shape[3];

More of a question, but does this relate to the complex64 change? No need to explain how if it is.

---

src/kernels/backend_webgl.ts, line 104 at r2 (raw file):

}

export interface TensorHandle {

Consider adding a comment for what a TensorHandle is

[nsthorat]

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

---

src/tensor.ts, line 38 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

remove this property

Done.

---

src/tensor.ts, line 82 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

no need to say as number

Done.

---

src/tensor.ts, line 389 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

remove

done

---

src/tensor.ts, line 412 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

undo this assert now that storageSize is gone.

removed because backend cpu and backend webgl are identical

---

src/tensor.ts, line 521 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

this is getting into tricky API territory, but right now it's strange that get() will return a single number instead of a tuple and the indexing is implicitly halved. File an issue to track this , but don't handle it here (it's a rare use-case to use .get() on complex)

tensorflow/tfjs#653

---

src/tensor_test.ts, line 1081 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

shouldn't this be 6 * BYTES_PER_COMPLEX, since at the end of the day, there are 12 floats (124 bytes), or 6 "complex floats" (68 bytes). Or I guess, what' happening is that numBytes doubles when registerTensor() gets called with the complex dtype. This is a failure of the current design and assumption of engine. Add a comment here why the number of bytes doubled. Useful when we revisit this test.

You are right. The problem was that backend cpu and backend webgl behaved differently. I consolidated the behavior so that they are identical (backend cpu holds onto real and imag tensors), and avoided counting bytes for complex64 tensors across the board.

---

src/util.ts, line 161 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

remove

this is no longer needed after we consolidated backend cpu and backend webgl

---

src/util.ts, line 338 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

undo the changes to this method. IIUC, you don't call this method with complex64 anymore.

Done.

---

src/kernels/backend_cpu.ts, line 184 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

backend_webl calls Tensor.make(real.shape, {}) and uses extra book-keeping to keep pointers to the real and img tensors. backend_cpu doesn't do this. Should we align them?

aligning made everything simpler, specifically w.r.t counting bytes. now we never count complex64 bytes, only counting the underlying tensors.

---

src/kernels/backend_cpu.ts, line 2335 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

More of a question, but does this relate to the complex64 change? No need to explain how if it is.

Nothing changed here except I'm asserting dy not complex

---

src/kernels/backend_webgl.ts, line 104 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Consider adding a comment for what a TensorHandle is

Done.

---

src/kernels/webgl/tex_util.ts, line 38 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

complexDataIds field doesn't exist. did you mean complexTensors fields?

Ha nice catch, done.

---

src/ops/arithmetic_test.ts, line 513 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Do you know why memory is different across cpu and webgl? Can we make those tests run on both cpu and webgl, but not node (easy to do by passing {IS_NODE: false} as test constraint).

It used to be different because the number of tensors is different (in WebGL we have 2 extra tensors). Now everything is aligned.

---

src/ops/array_ops_test.ts, line 3570 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

I feel like array_ops_test is not the right file for these tests. How about moving them to complex64_test.ts ?

Done.

---

src/ops/array_ops_test.ts, line 3633 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Add one more unit test where you make a real, img, and complex, and dispose the real and img, without disposing the complex, and making sure dataSync() still works on the complex. Smth like:

real = tensor(...)
img = tensor(...)
complex = tf.complex(real, img);
real.dispose()
img.dispose();
// Test memory
// Test complex.dataSync()
complex.dispose();
// Test memory

Done.

[dsmilkov]

Reviewed 4 of 10 files at r3.
Reviewable status: complete! 2 of 1 approvals obtained (waiting on @tafsiri and @nsthorat)