概述
本指南概述了在 TensorFlow.js 中定义自定义操作 (op)、内核和梯度的机制。它旨在概述主要概念并提供指向演示概念的代码的指针。
本指南适用于谁？
这是一份相当高级的指南，涉及 TensorFlow.js 的一些内部机制，它可能对以下人群特别有用
- 对 TensorFlow.js 感兴趣的进阶用户，他们希望自定义各种数学运算的行为（例如，研究人员覆盖现有梯度实现或需要修补库中缺失功能的用户）
- 构建扩展 TensorFlow.js 的库的用户（例如，在 TensorFlow.js 原语之上构建的通用线性代数库或新的 TensorFlow.js 后端）。
- 有兴趣为 tensorflow.js 贡献新操作的用户，他们希望了解这些机制的工作原理。
这 **不是** 一份关于 TensorFlow.js 通用使用的指南，因为它涉及内部实现机制。您无需了解这些机制即可使用 TensorFlow.js
您需要熟悉（或愿意尝试）阅读 TensorFlow.js 源代码才能充分利用本指南。
术语
对于本指南，提前描述一些关键术语很有用。
**操作 (Ops)** — 对一个或多个张量进行的数学运算，产生一个或多个张量作为输出。Ops 是“高级”代码，可以使用其他 Ops 来定义其逻辑。
**内核** — 与特定硬件/平台功能绑定的特定操作实现。内核是“低级”且特定于后端的。一些 Ops 从 Op 到内核有一对一的映射，而其他 Ops 使用多个内核。
**梯度/GradFunc** — **op/内核** 的“反向模式”定义，它计算该函数相对于某个输入的导数。梯度是“高级”代码（不是特定于后端的），可以调用其他 Ops 或内核。
**内核注册表** - 从 **(内核名称，后端名称)** 元组到内核实现的映射。
**梯度注册表** — 从 **内核名称到梯度实现** 的映射。
代码组织
内核特定于后端，并在它们各自的后端文件夹中定义（例如，tfjs-backend-cpu）。
自定义操作、内核和梯度不需要在这些包中定义。但它们在实现中通常会使用类似的符号。
实现自定义操作
自定义操作的一种思考方式是，它只是一个返回一些张量输出的 JavaScript 函数，通常以张量作为输入。
- 一些操作可以用现有的操作完全定义，并且应该直接导入和调用这些函数。 这是一个示例。
- 操作的实现也可以分派到特定于后端的内核。这是通过
Engine.runKernel完成的，将在“实现自定义内核”部分进一步描述。 这是一个示例。
实现自定义内核
特定于后端的内核实现允许针对给定操作的逻辑进行优化实现。内核由调用 tf.engine().runKernel() 的操作调用。内核实现由四件事定义
- 内核名称。
- 实现内核的后端。
- 输入：内核函数的张量参数。
- 属性：内核函数的非张量参数。
这是一个 内核实现示例。用于实现的约定特定于后端，最好从查看每个特定后端的实现和文档来理解。
通常，内核在比张量更低的级别上运行，而是直接读写最终将由 tfjs-core 包装成张量的内存。
实现内核后，可以使用 tfjs-core 中的 registerKernel 函数 将其注册到 TensorFlow.js。您可以为要使用该内核的每个后端注册一个内核。注册后，可以使用 tf.engine().runKernel(...) 调用内核，TensorFlow.js 将确保分派到当前活动后端中的实现。
实现自定义梯度
梯度通常针对给定内核定义（由在调用 tf.engine().runKernel(...) 时使用的相同内核名称标识）。这使 tfjs-core 能够使用注册表在运行时查找任何内核的梯度定义。
实现自定义梯度对于以下情况很有用
- 添加库中可能不存在的梯度定义
- 覆盖现有的梯度定义，以自定义给定内核的梯度计算。
您可以在此处查看梯度实现示例。
实现给定调用的梯度后，可以使用registerGradient 函数（来自 tfjs-core）将其注册到 TensorFlow.js。
实现自定义梯度的另一种方法是使用tf.customGrad，它绕过梯度注册表（因此允许以任意方式为任意函数计算梯度）。
以下是在库中使用 customGrad 的操作示例