Modularize the remaining unary ops in cpu backend

[jinjingforever]

Most ops are fairly simple, taking advantage of newly created unaryKernelFunc helper function.

---

This change is 

[tafsiri]

Great stuff, left a couple of small comments on documentation and a larger one on supporting complex outputs from unary ops.

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

---

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

export const clipKernelFunc = unaryKernelFunc(ClipByValue, (x, attrs) => {
  const clipAttrs = attrs as {} as ClipByValueAttrs;
  return x > clipAttrs.clipValueMax ?

lets use an if statement here to avoid nesting ternaries below.

---

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

import {DataType, NamedAttrMap, TypedArray} from '@tensorflow/tfjs-core';

As we add more of these I think its worth adding a comment explaining the naming convention, (i.e. ...Operation vs ...KernelImpl), looking at this diff had me wondering why there isn't a SimpleUnaryKernelImpl

---

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

export type SimpleUnaryOperation = (x: number, attrs?: NamedAttrMap) => number;
export type ComplexUnaryOperation =
    (real: number, imag: number, attrs?: NamedAttrMap) => number;

This suggests all complex unary operations will produce real numbers (as opposed to producing complex numbers) right? I think in the abs example you have in this PR that is true, but I do believe other kernels would define complex results for complex inputs. I tried the following in a python colab

real_part = tf.random.uniform([3,4,5])
input_tensor = tf.complex(real=real_part, imag=tf.zeros_like(real_part))
print(input_tensor.dtype)
res = tf.sin(input_tensor)
print(res.dtype)
res

and got the following result

<dtype: 'complex64'>
<dtype: 'complex64'>
<tf.Tensor: shape=(3, 4, 5), dtype=complex64, numpy=
array([[[0.38650748+0.j, 0.533613  +0.j, 0.62358844+0.j, 0.35709518+0.j,
         0.11534068+0.j],
  ... (results continue)

I think we want to leave room for supporting this kind of op as well

---

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

 * @param op A `SimpleUnaryOperation` for the kernel.
 * @param dtype Optional. If set, the result has this dtype. Otherwise, the
 *     result has the same dtype as the first input. This is mainly used in

nit: same as the input (remove first since there is only one input)

---

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

    if (x.dtype === 'complex64') {
      util.assert(
          opComplex !== undefined, () => `no complex op defined for ${name}`);

use != null to cover both null and undefined.

---

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

      }
    }
    return cpuBackend.makeTensorInfo(x.shape, dtype || x.dtype, newValues);

If the result of the operation is supposed to be a complex tensor I don't think this will construct it in the right way (i think it would need to go through the complex kernel util), this is related to my comment above on the ComplexUnaryOperation type.

[lina128]

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

---

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

export const absKernelFunc = unaryKernelFunc(
    Abs, (x) => Math.abs(x), 'float32', (real, imag) => Math.hypot(real, imag));

I suggest we don't use unaryKernelFunc for abs, because complexAbs seems not return a complex and we don't know whether this can be generalized to other unary ops.

---

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

export const absKernelFunc = unaryKernelFunc(
    Abs, (x) => Math.abs(x), 'float32', (real, imag) => Math.hypot(real, imag));

Can we change all x to xi for the passed in function? It denotes that the operation will apply on each element of x. Otherwise, it is confusing, given x represents tensor/tensorInfo most of the time.

---

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

import {DataType, NamedAttrMap, TypedArray} from '@tensorflow/tfjs-core';

export type SimpleUnaryOperation = (x: number, attrs?: NamedAttrMap) => number;

Should these be in a separate file? unary_types.ts

---

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

Previously, tafsiri (Yannick Assogba) wrote…

This suggests all complex unary operations will produce real numbers (as opposed to producing complex numbers) right? I think in the abs example you have in this PR that is true, but I do believe other kernels would define complex results for complex inputs. I tried the following in a python colab

real_part = tf.random.uniform([3,4,5])
input_tensor = tf.complex(real=real_part, imag=tf.zeros_like(real_part))
print(input_tensor.dtype)
res = tf.sin(input_tensor)
print(res.dtype)
res

and got the following result

<dtype: 'complex64'>
<dtype: 'complex64'>
<tf.Tensor: shape=(3, 4, 5), dtype=complex64, numpy=
array([[[0.38650748+0.j, 0.533613  +0.j, 0.62358844+0.j, 0.35709518+0.j,
         0.11534068+0.j],
  ... (results continue)

I think we want to leave room for supporting this kind of op as well

+1

---

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

    const cpuBackend = backend as MathBackendCPU;

    const values = cpuBackend.readSync(x.dataId) as TypedArray;

This generates new array if type is complex64, an expensive step. I wonder whether we can refactor to just iterate through the underlying real and imag.

---

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

    const xSize = util.sizeFromShape(x.shape);
    const newValues =
        dtype === 'bool' ? new Uint8Array(xSize) : new Float32Array(xSize);

Can we use util.getArrayFromDType? Also what if dtype is null but x.dtype is 'bool'?

---

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

      }
    } else {
      assertNotComplex(x, name);

Do we need this assertion?

---

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

Previously, tafsiri (Yannick Assogba) wrote…

If the result of the operation is supposed to be a complex tensor I don't think this will construct it in the right way (i think it would need to go through the complex kernel util), this is related to my comment above on the ComplexUnaryOperation type.

If complexAbs is the only one that deal with complex64, I suggest we don't generalize it for now. We can just leave opComplex out.

[jinjingforever]

Thanks Yannick and Na! As suggested, I removed the logic of handling complex numbers from unaryKernelFunc. I also update Abs.ts to NOT use unaryKernelFunc, and refactored the code to operate on the underlying data directly (instead of using readSync and merge real and imag values into a single new array).

PTAL

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

---

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

Previously, lina128 (Na Li) wrote…

I suggest we don't use unaryKernelFunc for abs, because complexAbs seems not return a complex and we don't know whether this can be generalized to other unary ops.

Done

---

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

Previously, lina128 (Na Li) wrote…

Can we change all x to xi for the passed in function? It denotes that the operation will apply on each element of x. Otherwise, it is confusing, given x represents tensor/tensorInfo most of the time.

Done

---

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

Previously, tafsiri (Yannick Assogba) wrote…

lets use an if statement here to avoid nesting ternaries below.

Done.

---

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

Previously, tafsiri (Yannick Assogba) wrote…

As we add more of these I think its worth adding a comment explaining the naming convention, (i.e. ...Operation vs ...KernelImpl), looking at this diff had me wondering why there isn't a SimpleUnaryKernelImpl

Good point! Now I moved SimpleUnaryOperation to unary_types.ts. Will add more comments when it becomes more complex.

---

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

Previously, lina128 (Na Li) wrote…

Should these be in a separate file? unary_types.ts

Good catch! Done

---

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

Previously, lina128 (Na Li) wrote…

+1

Ack. Removed the generic logic of handling complex numberfs. Will add as needed in the future

---

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

Previously, tafsiri (Yannick Assogba) wrote…

nit: same as the input (remove first since there is only one input)

Done

---

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

Previously, lina128 (Na Li) wrote…

This generates new array if type is complex64, an expensive step. I wonder whether we can refactor to just iterate through the underlying real and imag.

Done

---

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

Previously, lina128 (Na Li) wrote…

Can we use util.getArrayFromDType? Also what if dtype is null but x.dtype is 'bool'?

Done. Introduced $dtype which is "dtype || x.dtype" to make types consistent. Also added a check to make sure we are not operating on strings.

---

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

Previously, tafsiri (Yannick Assogba) wrote…

use != null to cover both null and undefined.

Ack (this part is removed)

---

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

Previously, lina128 (Na Li) wrote…

Do we need this assertion?

Ack (this part is removed. Keeping this to make sure inputs are not complex numbers)

---

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

Previously, lina128 (Na Li) wrote…

If complexAbs is the only one that deal with complex64, I suggest we don't generalize it for now. We can just leave opComplex out.

Done. Removed the generic logic of handling complex numbers. Will add as needed in the future.

[lina128]

Great work, this PR is very clean and really expedite the modularization effort, thank you Jing!!!

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

---

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

 *     inputs.
 */
export function unaryKernelFunc(

Let's seperate this into another file, unary_utils.ts, because this file is going to be renamed to binary_utils.ts.

[tafsiri]

👏

Reviewed 39 of 39 files at r2.
Reviewable status: complete! 2 of 1 approvals obtained (waiting on @jinjingforever, @lina128, and @tafsiri)