Remove kernel_registry and call backends directly.

[dsmilkov]

• Remove kernel_registry and call backends directly
• Allow for forward function to save tensors that will be needed by the backend function. We pass dy and the saved tensors to the backprop function.
• Remove logic in tape to handle multiple outputs. We can bring it back if needed.
• consolidate tape_types.ts and tape_util.ts into a single file tape.ts
• in debug mode, log the name of the activeScope, which should point to the last operation which called the kernel. The aggregation of these and showing hierarchical debug view is still a TODO.

This change is 

[nsthorat]

Awesome work, glad you're killing the tape types :)

Only major comment is in tensor.ts

---

Reviewed 49 of 79 files at r3.
Review status: 49 of 58 files reviewed at latest revision, 3 unresolved discussions, some commit checks failed.

---

src/engine.ts, line 37 at r3 (raw file):

export type ForwardFunc<T extends Tensor> =
    (backend: KernelBackend, save?: <S extends Tensor>(tensor: S) => S) => T;

it'd be nice to document "save" somewhere, maybe here. Just say that these are tensors computed in the forward pass that we need in the backwards pass.

---

src/tape.ts, line 121 at r3 (raw file):

      const prunedNode = Object.assign({}, node) as TapeNode;
      prunedNode.inputs = prunedInputs;
      prunedNode.output = prunedOutputs;

you can consolidate this block to:

prunedNode.output = node.output

and completely remove "prunedOutputs"

---

src/tensor.ts, line 598 at r3 (raw file):

    return ops.subStrict(this, x);
  }
  pow<T extends Tensor>(this: T, exp: Tensor): T {

this change seems a little weird to me - why are the chained ops now taking this as the first argument, and why not all of them?

---

Comments from Reviewable

[nsthorat]

Review status: 49 of 58 files reviewed at latest revision, 4 unresolved discussions, some commit checks failed.

---

src/engine.ts, line 102 at r3 (raw file):

      return x;
    };
    const kernelName = this.activeScope.name;

this is kind of funny, won't the active scope name be the operation name because of @operation and the scope around the operation, and not the actual kernelName?

---

Comments from Reviewable

[nsthorat]

Review status: 49 of 58 files reviewed at latest revision, 4 unresolved discussions, some commit checks failed.

---

src/engine.ts, line 102 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

this is kind of funny, won't the active scope name be the operation name because of @operation and the scope around the operation, and not the actual kernelName?

One idea is to keep passing a string for the kernel name just for logging.

---

Comments from Reviewable

[dsmilkov]

Review status: 48 of 58 files reviewed at latest revision, 4 unresolved discussions, some commit checks failed.

---

src/engine.ts, line 37 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

it'd be nice to document "save" somewhere, maybe here. Just say that these are tensors computed in the forward pass that we need in the backwards pass.

Done.

---

src/engine.ts, line 102 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

One idea is to keep passing a string for the kernel name just for logging.

Great observation. So I was thinking about this a little.

There will be a skew between kernels and user-ops and I think for now, before we have aggregation done, it's better to report user-ops since that is what the user controls (we report the inner-most named scope, with ability to aggregate).

Couple reasons for this argument:

• When we backpropagate, if we see that a tapeNode is missing gradFunc, (meaning we haven't implemented gradient yet), it's more useful to report an error that this opName is missing gradients, so we can't backpropagate, as opposed to this kernelName is missing grad, since from the user perspective, the culprit is dl.opName(), not dl.kernelName().
• When reporting timing in debug mode, kernelNames might be confusing, e.g. imagine we reuse a conv2d kernel for higher-order convs, reporting the logical conv is more helpful than reporting conv2d (from the user perspective).

Once we have aggregation done, we can come back and decide maybe to report both kernel and op, at which point, we can add a required name argument to engine.executeKernel().

WDYT?

---

src/tape.ts, line 121 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

you can consolidate this block to:

prunedNode.output = node.output

and completely remove "prunedOutputs"

Done.

---

src/tensor.ts, line 598 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

this change seems a little weird to me - why are the chained ops now taking this as the first argument, and why not all of them?

this: T is not a real argument, it's a special signal to the TS compiler that if this is of type T, when this.pow() will return the same type. Previously we had no constraint that a.pow() will return the same rank as a. Now we do.

---

Comments from Reviewable

[nsthorat]

Review status: 46 of 58 files reviewed at latest revision, 4 unresolved discussions, all commit checks successful.

---

src/engine.ts, line 102 at r3 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Great observation. So I was thinking about this a little.

There will be a skew between kernels and user-ops and I think for now, before we have aggregation done, it's better to report user-ops since that is what the user controls (we report the inner-most named scope, with ability to aggregate).

Couple reasons for this argument:

• When we backpropagate, if we see that a tapeNode is missing gradFunc, (meaning we haven't implemented gradient yet), it's more useful to report an error that this opName is missing gradients, so we can't backpropagate, as opposed to this kernelName is missing grad, since from the user perspective, the culprit is dl.opName(), not dl.kernelName().
• When reporting timing in debug mode, kernelNames might be confusing, e.g. imagine we reuse a conv2d kernel for higher-order convs, reporting the logical conv is more helpful than reporting conv2d (from the user perspective).

Once we have aggregation done, we can come back and decide maybe to report both kernel and op, at which point, we can add a required name argument to engine.executeKernel().

WDYT?

That sounds reasonable to me. Maybe let's name this field to that effect: s/kernelName/scopeName or something.

---

Comments from Reviewable

[dsmilkov]

Review status: 46 of 58 files reviewed at latest revision, 4 unresolved discussions, all commit checks successful.

---

src/engine.ts, line 102 at r3 (raw file):

Previously, nsthorat (Nikhil Thorat) wrote…

That sounds reasonable to me. Maybe let's name this field to that effect: s/kernelName/scopeName or something.

Done.

---

Comments from Reviewable