[core]Modularize addN.

[lina128]

Changes in this PR include:
(1) Split AddN out from binary_ops.
(2) Split AddN's gradient out to a separate file.
(3) Add logic to engine to handle array of inputs to save.

Context for (3), the gradient modularization requires a gradConfig for each gradient function in their separate file. This gradConfig has a inputsToSave field, currently the logic is that inputsToSave field will specify tensors to save: https://github.com/tensorflow/tfjs/blob/master/tfjs-core/src/engine.ts#L669. But it couldn't handle the case that the inputs have dynamic length of Tensors to save. To deal with this, we can introduce a flag saveAllInputs to gradConfig, which indicates to engine that simply save all the inputs.

An alternative is to change the inputsToSave type from Tensor[] to Array<Tensor|Tensor[]>, but this generates more issues than desired, for example, a lot of the places (logging, profiling, backward functions, etc.) expect the simple form Tensor[], and has logic to access each array element and their shape or dtype, etc. Changing will cause to add more logic to evaluate whether it is an array or a tensor. See this PR for the places to change, and this PR still fails with 10 more places to fix, at which point, I gave up on this solution: https://github.com/tensorflow/tfjs/compare/master...lina128:addn_try?expand=1

I think the first solution is much simpler and elegant.

---

This change is 

[lina128]

We need to have an exception here, AddNInputs is an array of Tensors. We don't have a fixed number of Tensors, so we cannot use NamedTensorInfoMap. Cast this input as NamedTensorMap and use it later is fine, because array is basically an object with index as key, this array can use all the Object apis.

[tafsiri]

Left a set of comments around exploring a different more explicit design. Let me know if you want to chat more about this.

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

---

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

      let inputTensorsToSave: Tensor[];
      if (!inputsToSave) {

See comment below on NamedTensorInfoMap. I think this design might get confusing when there are multiple inputs and some might be arrays. The move to a map lets be more specific on which derivatives map to which inputs. We can chat via GVC if you want to discuss this alternative more.

---

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

Previously, lina128 (Na Li) wrote…

We need to have an exception here, AddNInputs is an array of Tensors. We don't have a fixed number of Tensors, so we cannot use NamedTensorInfoMap. Cast this input as NamedTensorMap and use it later is fine, because array is basically an object with index as key, this array can use all the Object apis.

I thought we discussed changing NamedTensorInfoMap to allow the value to be an array of TensorInfo's as well. I think that would be clearner/better express the intent

---

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

export interface GradConfig {
  kernelName: string;
  // If the array is empty, all the inputs will be saved.

We do not want to save the inputs unless the gradient actually requests it. I think specifying empty array vs undefined is confusing.

---

tfjs-core/src/gradients/AddN_grad.ts, line 24 at r1 (raw file):

export const addNGradConfig: GradConfig = {
  kernelName: AddN,
  inputsToSave: [],

See comments above about allowing NamedTensorInfoMap to go from String => TensorInfo|TensorInfo[]

[lina128]

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

---

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

Previously, tafsiri (Yannick Assogba) wrote…

See comment below on NamedTensorInfoMap. I think this design might get confusing when there are multiple inputs and some might be arrays. The move to a map lets be more specific on which derivatives map to which inputs. We can chat via GVC if you want to discuss this alternative more.

Changed to use a flag to indicate whether save all inputs. Thanks Yannick!

---

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

Previously, tafsiri (Yannick Assogba) wrote…

I thought we discussed changing NamedTensorInfoMap to allow the value to be an array of TensorInfo's as well. I think that would be clearner/better express the intent

That solution caused too much change in logic, see example PR: https://github.com/tensorflow/tfjs/compare/master...lina128:addn_try?expand=1

---

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

Previously, tafsiri (Yannick Assogba) wrote…

We do not want to save the inputs unless the gradient actually requests it. I think specifying empty array vs undefined is confusing.

Agreed. Change to use a flag.

[lina128]

Discussed offline with @tafsiri, the first solution (the one in this PR) works with current code with minimal change. We just need to introduce a flag to gradConfig to handle the case where inputs has a dynamic number of elements to save and we just want to save all inputs. The second solution (in another PR, link in this PR's description) requires logic change in many places, need to understand from Daniel and Nikhil whether this is needed. And what's the original thoughts to handle array of Tensors as object of Tensors with array index as key.

[tafsiri]

Thanks Na, updated the review. A few changes (adding a test and some assertions). Thanks!

Reviewed 5 of 5 files at r2.
Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, @lina128, @nsthorat, and @tafsiri)

---

scripts/release-notes.ts, line 0 at r2 (raw file):
This may have been committed by accident?

---

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

Previously, lina128 (Na Li) wrote…

Changed to use a flag to indicate whether save all inputs. Thanks Yannick!

Great! Just to recap for anyone else reviewing, We discussed offline and after she walked me through the alternatives agreed with Na that this solution is fine. We have very few ops that take lists of tensors and none that take tensor lists and regular tensors.

---

tfjs-core/src/engine.ts, line 674 at r2 (raw file):

      // If saveAllInputs is true, all inputs will be saved. Otherwise, inputs
      // specified in inputsToSave will be saved.
      const inputTensorsToSave: Tensor[] = gradConfig.saveAllInputs ?

When saveAllInputs is true assert that inputs is an array before doing this mapping.

Else I think you need to handle the case where inputs is a map where values might be arrays or tensors (at which point you would have to check there is only one array to avoid key collision). Which we could do if we run into that situation.

While our type system prevents this. We would want to signal this early to anyone trying to implement custom ops and gradients outside of typescript.

---

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

Previously, lina128 (Na Li) wrote…

That solution caused too much change in logic, see example PR: https://github.com/tensorflow/tfjs/compare/master...lina128:addn_try?expand=1

Chatted offline, updates above. I did take a quick look into TF to see how they handle it and they do have a DT_VARIANT type. So this could be an analog to that.

---

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

Previously, tafsiri (Yannick Assogba) wrote…

We do not want to save the inputs unless the gradient actually requests it. I think specifying empty array vs undefined is confusing.

After talking offline decided to use a flag to indicate saving all inputs.

---

tfjs-core/src/kernel_registry.ts, line 64 at r2 (raw file):

  inputsToSave?: string[];
  // If saveAllInputs is true, inputsToSave will be ignored, all inputs will
  // be saved.

Depending on how you resolve the above comment in engine, Add to this comment to indicate that it supports only passing in a single array of tensors as input.

---

tfjs-core/src/kernel_registry.ts, line 65 at r2 (raw file):

  // If saveAllInputs is true, inputsToSave will be ignored, all inputs will
  // be saved.
  saveAllInputs?: boolean;

Could you add a test that exercises this in kernel_registry_test? There is a section for registering gradients.

[lina128]

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

---

scripts/release-notes.ts, line at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

This may have been committed by accident?

Removed. Nice catch!

---

tfjs-core/src/engine.ts, line 674 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

When saveAllInputs is true assert that inputs is an array before doing this mapping.

Else I think you need to handle the case where inputs is a map where values might be arrays or tensors (at which point you would have to check there is only one array to avoid key collision). Which we could do if we run into that situation.

While our type system prevents this. We would want to signal this early to anyone trying to implement custom ops and gradients outside of typescript.

Hi Yannick, can you give an example? I can't think of a case that will have key collision. If there is an element in the map that is an array, it will be one element as an array in the inputTensorsToSave, not N elements. Essentially, we are not changing the way current code works, Object.keys(inputs).map((key) => inputs[key]) has the same effect as inputsToSave.map((inputName) => inputs[inputName]), except the former copies every element to the result, whereas the later selectively copies element to the result.

---

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

Previously, tafsiri (Yannick Assogba) wrote…

Chatted offline, updates above. I did take a quick look into TF to see how they handle it and they do have a DT_VARIANT type. So this could be an analog to that.

Good to know!

---

tfjs-core/src/kernel_registry.ts, line 65 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Could you add a test that exercises this in kernel_registry_test? There is a section for registering gradients.

Great suggestion, done!

[tafsiri]

LGTM. Good work on this. Added more description of the safety check I think we should make.

Reviewed 2 of 2 files at r3.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan, @dsmilkov, @lina128, and @nsthorat)

---

tfjs-core/src/engine.ts, line 674 at r2 (raw file):

Previously, lina128 (Na Li) wrote…

Hi Yannick, can you give an example? I can't think of a case that will have key collision. If there is an element in the map that is an array, it will be one element as an array in the inputTensorsToSave, not N elements. Essentially, we are not changing the way current code works, Object.keys(inputs).map((key) => inputs[key]) has the same effect as inputsToSave.map((inputName) => inputs[inputName]), except the former copies every element to the result, whereas the later selectively copies element to the result.

Sure. I'm thinking of someone constructing an input like

const inputs = {
  a: [tf.zeros([1,2]), tf.zeros([2,1])], 
  b: [tf.zeros([3,1]), tf.zeros([1,3])]
};

And hoping that saveAllInputs will flatten the arrays in a and b. While typescript prevents us from doing this, one could do this in JS. Asserting that inputs is an array would just reject this and could provide an easier to understand method. I believ having an element of inputTensorsToSave be an array rather than a Tensor would break the other parts of engine and tape that expect that to be a tensor.

While this was also true in the old saveFunc, I'm thinking that since we are adding this explicit saveAllInputs flag, there is a greater chance of misinterpreting what it can do. It is essentially designed for a flat array of tensors as input.

Separately (and optionally) I think its okay to disallow using saveAllInputs with actual named input maps. I can't quite think of a scenario where there is a variable number of inputs that are in a map that we wouldn't want to spell out in inputsToSave. And in our code we should not rely on that as a shorthand for not specifying the inputs by name.

---

tfjs-core/src/kernel_registry_test.ts, line 354 at r3 (raw file):

    expect(dx[1].dtype).toBe('float32');
    expect(dx[1].shape).toEqual([2, 2]);
    // expectArraysClose(await dx[0].data(), [3, 3, 3, 3]);

remove commented out code.

[lina128]

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

---

tfjs-core/src/engine.ts, line 674 at r2 (raw file):

Previously, tafsiri (Yannick Assogba) wrote…

Sure. I'm thinking of someone constructing an input like

const inputs = {
  a: [tf.zeros([1,2]), tf.zeros([2,1])], 
  b: [tf.zeros([3,1]), tf.zeros([1,3])]
};

And hoping that saveAllInputs will flatten the arrays in a and b. While typescript prevents us from doing this, one could do this in JS. Asserting that inputs is an array would just reject this and could provide an easier to understand method. I believ having an element of inputTensorsToSave be an array rather than a Tensor would break the other parts of engine and tape that expect that to be a tensor.

While this was also true in the old saveFunc, I'm thinking that since we are adding this explicit saveAllInputs flag, there is a greater chance of misinterpreting what it can do. It is essentially designed for a flat array of tensors as input.

Separately (and optionally) I think its okay to disallow using saveAllInputs with actual named input maps. I can't quite think of a scenario where there is a variable number of inputs that are in a map that we wouldn't want to spell out in inputsToSave. And in our code we should not rely on that as a shorthand for not specifying the inputs by name.

Understand. I was originally thinking the flag be a convenient way to specify inputs for both array and map (optional). But yeah, maybe it's a good thing to make the use case narrower and throw error for cases we didn't expect. This way, in the future if such case surface, we would know. Changed the logic so that the flag should only be used with an array of Tensors. Also changed the test to check if code throw error if flag and map is used together.