backend: Add ops onesLike and zerosLike

[caisq]

• Also add CPU and WebGL backend implementations of these
  based on fill()
• This will be used in tfjs-node to bind to the libtensorflow
  implementations of onesLike and zerosLike to save memory
  copying between CPU and GPU, which will lead to performance
  improvements.

PERF

---

This change is 

[caisq]

cc @Lewuathe

[dsmilkov]

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

---

src/kernels/backend.ts, line 592 at r2 (raw file):

  }

  onesLike<R extends Rank>(x: Tensor<R>): Tensor<R> {

adding them to the the backend is not enough. You have to edit onesLike_ and zerosLike_ in tensor_ops.ts to call ENV.engine.runKernel(backend => backend.onesLike(...)) in order to delegate user code to the backend.

---

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

  onesLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
    if (x.dtype === 'string') {
      throw new Error('onesLike is not supported under string dtype');

nit: for tensors of string dtype

---

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

  zerosLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
    return this.fill(x.shape, x.dtype === 'string' ? '' : 0, x.dtype);

calling fill is sub-optimal here since newly created typedarrays are already filled with zeros so you want to avoid calling typedarray.fill(0). Call Tensor.make(shape, {values}, dtype), where values = getArrayFromDType(...)

---

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

      throw new Error('onesLike is not supported under string dtype');
    } else {
      return this.fill(x.shape, 1, x.dtype);

let's add a gpu shader for these.

[dsmilkov]

Reviewable status: 0 of 1 approvals obtained (waiting on @caisq and @dsmilkov)

---

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

  zerosLike<R extends Rank>(x: Tensor<R>): Tensor<R> {
    return this.fill(x.shape, x.dtype === 'string' ? '' : 0, x.dtype);

Does TF python assign an empty string for tf.zeros_like() ?

[caisq]

Thanks for the review!

Reviewable status: 0 of 1 approvals obtained (waiting on @caisq and @dsmilkov)

---

src/kernels/backend.ts, line 592 at r2 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

adding them to the the backend is not enough. You have to edit onesLike_ and zerosLike_ in tensor_ops.ts to call ENV.engine.runKernel(backend => backend.onesLike(...)) in order to delegate user code to the backend.

Done. Sorry I had this in mind but forgot to push the commit to the PR.

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

nit: for tensors of string dtype

Done with minor modification.

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

Does TF python assign an empty string for tf.zeros_like() ?

Yes. Confirmed that. But it appears that the HEAD of the master branch of tfjs-core doesn't support it yet. (There is an assertDType call by zerosLike that throws an error on string dtype). We can fix that in a later PR.

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

calling fill is sub-optimal here since newly created typedarrays are already filled with zeros so you want to avoid calling typedarray.fill(0). Call Tensor.make(shape, {values}, dtype), where values = getArrayFromDType(...)

Done.

[dsmilkov]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @caisq)

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

let's add a gpu shader for these.

Want to add a todo and file an issue to turn these as webgl shaders?

[caisq]

Reviewable status: complete! 1 of 1 approvals obtained

---

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

Previously, dsmilkov (Daniel Smilkov) wrote…

Want to add a todo and file an issue to turn these as webgl shaders?

Done.

[dsmilkov]

Reviewed 3 of 3 files at r3.
Reviewable status: complete! 1 of 1 approvals obtained