建立運算 (op)

如果您想建立現有 TensorFlow 程式庫未涵蓋的運算 (op)，建議您先嘗試以 Python 撰寫運算 (op)，作為現有 Python 運算 (ops) 或函式的組合。如果不可行，您可以建立自訂 C++ 運算 (op)。您可能想要建立自訂 C++ 運算 (op) 的原因有幾個：

• 不容易或無法將您的運算表示為現有運算 (ops) 的組合。
• 將您的運算表示為現有基本元素的組合效率不高。
• 您想要手動融合基本元素的組合，而未來的編譯器會發現難以融合。

例如，假設您想要實作類似「中位數池化」的功能，與「MaxPool」運算子類似，但計算滑動視窗中的中位數，而不是最大值。使用運算組合來執行此操作可能可行 (例如，使用 ExtractImagePatches 和 TopK)，但效能或記憶體效率可能不如原生運算，在原生運算中，您可以在單一融合運算中執行更聰明的操作。一如既往，通常首先值得嘗試使用運算子組合來表達您想要的功能，只有在證明困難或效率不高時，才選擇新增運算。

若要納入您的自訂運算 (op)，您需要：

1. 在 C++ 檔案中註冊新的運算 (op)。運算 (op) 註冊會為運算 (op) 的功能定義介面 (規格)，這與運算 (op) 的實作無關。例如，運算 (op) 註冊會定義運算 (op) 的名稱和運算 (op) 的輸入和輸出。它也會定義用於張量形狀推論的形狀函式。
2. 在 C++ 中實作運算 (op)。運算 (op) 的實作稱為核心 (kernel)，它是您在步驟 1 中註冊的規格的具體實作。針對不同的輸入/輸出類型或架構 (例如，CPU、GPU)，可以有多個核心 (kernels)。
3. 建立 Python 封裝函式 (wrapper) (選用)。此封裝函式 (wrapper) 是用於在 Python 中建立運算 (op) 的公開 API。預設封裝函式 (wrapper) 是從運算 (op) 註冊產生的，可以直接使用或新增。
4. 編寫函式以計算運算 (op) 的梯度 (選用)。
5. 測試運算 (op)。為了方便起見，我們通常在 Python 中執行此操作，但您也可以在 C++ 中測試運算 (op)。如果您定義梯度，可以使用 Python tf.test.compute_gradient_error 驗證梯度。請參閱 relu_op_test.py，其中範例測試了 Relu 類運算子的前向函式及其梯度。

事前準備

• 對 C++ 有一定程度的熟悉。
• 必須已安裝 TensorFlow 二進位檔，或必須已下載 TensorFlow 原始碼，並且能夠建構。

定義運算 (op) 介面

您可以透過向 TensorFlow 系統註冊來定義運算 (op) 的介面。在註冊中，您可以指定運算 (op) 的名稱、其輸入 (類型和名稱) 和輸出 (類型和名稱)，以及文件字串和運算 (op) 可能需要的任何屬性 (attrs)。

若要瞭解其運作方式，假設您想要建立一個運算 (op)，該運算 (op) 接受 int32 張量，並輸出張量的副本，但第一個元素以外的所有元素都設為零。若要執行此操作，請建立名為 zero_out.cc 的檔案。然後新增對 REGISTER_OP 巨集的呼叫，以定義運算 (op) 的介面

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"

using namespace tensorflow;

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

此 ZeroOut 運算 (op) 接受一個 32 位元整數張量 to_zero 作為輸入，並輸出一個 32 位元整數張量 zeroed。此運算 (op) 也使用形狀函式來確保輸出張量的形狀與輸入張量的形狀相同。例如，如果輸入是形狀為 [10, 20] 的張量，則此形狀函式會指定輸出形狀也為 [10, 20]。

實作運算 (op) 的核心 (kernel)

定義介面後，請提供運算 (op) 的一或多個實作。若要建立其中一個核心 (kernel)，請建立一個擴充 OpKernel 並覆寫 Compute 方法的類別。Compute 方法提供一個 context 引數，其類型為 OpKernelContext*，您可以從中存取輸入和輸出張量等實用項目。

將您的核心 (kernel) 新增至您在上面建立的檔案。核心 (kernel) 可能看起來像這樣

#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<int32>();

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));
    auto output_flat = output_tensor->flat<int32>();

    // Set all but the first element of the output tensor to 0.
    const int N = input.size();
    for (int i = 1; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value if possible.
    if (N > 0) output_flat(0) = input(0);
  }
};

實作核心 (kernel) 後，您需要向 TensorFlow 系統註冊。在註冊中，您可以指定此核心 (kernel) 將在哪些不同的限制下執行。例如，您可能有一個為 CPU 建立的核心 (kernel)，以及一個為 GPU 建立的獨立核心 (kernel)。

若要為 ZeroOut 運算 (op) 執行此操作，請將以下內容新增至 zero_out.cc

REGISTER_KERNEL_BUILDER(Name("ZeroOut").Device(DEVICE_CPU), ZeroOutOp);

多執行緒 CPU 核心 (kernels)

若要編寫多執行緒 CPU 核心 (kernel)，可以使用 work_sharder.h 中的 Shard 函式。此函式會跨執行緒將運算函式分片，這些執行緒已設定為用於運算內執行緒處理 (請參閱 config.proto 中的 intra_op_parallelism_threads)。

GPU 核心 (kernels)

GPU 核心 (kernel) 分為兩個部分實作：OpKernel 以及 CUDA 核心 (kernel) 及其啟動程式碼。

有時，OpKernel 實作在 CPU 和 GPU 核心 (kernel) 之間是通用的，例如在檢查輸入和配置輸出方面。在這種情況下，建議的實作方式是

1. 定義以裝置和張量的基本類型為範本的 OpKernel。
2. 為了執行輸出的實際運算，Compute 函式會呼叫範本化的函子結構。
3. CPUDevice 的函子特化是在同一個檔案中定義的，但 GPUDevice 的特化是在 .cu.cc 檔案中定義的，因為它將使用 CUDA 編譯器進行編譯。

以下是一個實作範例。

// kernel_example.h
#ifndef KERNEL_EXAMPLE_H_
#define KERNEL_EXAMPLE_H_

#include <unsupported/Eigen/CXX11/Tensor>

template <typename Device, typename T>
struct ExampleFunctor {
  void operator()(const Device& d, int size, const T* in, T* out);
};

#if GOOGLE_CUDA
// Partially specialize functor for GpuDevice.
template <typename T>
struct ExampleFunctor<Eigen::GpuDevice, T> {
  void operator()(const Eigen::GpuDevice& d, int size, const T* in, T* out);
};
#endif

#endif KERNEL_EXAMPLE_H_

// kernel_example.cc
#include "kernel_example.h"

#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/shape_inference.h"
#include "tensorflow/core/framework/op_kernel.h"

using namespace tensorflow;

using CPUDevice = Eigen::ThreadPoolDevice;
using GPUDevice = Eigen::GpuDevice;

REGISTER_OP("Example")
    .Attr("T: numbertype")
    .Input("input: T")
    .Output("input_times_two: T")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

// CPU specialization of actual computation.
template <typename T>
struct ExampleFunctor<CPUDevice, T> {
  void operator()(const CPUDevice& d, int size, const T* in, T* out) {
    for (int i = 0; i < size; ++i) {
      out[i] = 2 * in[i];
    }
  }
};

// OpKernel definition.
// template parameter <T> is the datatype of the tensors.
template <typename Device, typename T>
class ExampleOp : public OpKernel {
 public:
  explicit ExampleOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    // Create an output tensor
    Tensor* output_tensor = NULL;
    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
                                                     &output_tensor));

    // Do the computation.
    OP_REQUIRES(context, input_tensor.NumElements() <= tensorflow::kint32max,
                errors::InvalidArgument("Too many elements in tensor"));
    ExampleFunctor<Device, T>()(
        context->eigen_device<Device>(),
        static_cast<int>(input_tensor.NumElements()),
        input_tensor.flat<T>().data(),
        output_tensor->flat<T>().data());
  }
};

// Register the CPU kernels.
#define REGISTER_CPU(T)                                          \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
      ExampleOp<CPUDevice, T>);
REGISTER_CPU(float);
REGISTER_CPU(int32);

// Register the GPU kernels.
#ifdef GOOGLE_CUDA
#define REGISTER_GPU(T)                                          \
  /* Declare explicit instantiations in kernel_example.cu.cc. */ \
  extern template class ExampleFunctor<GPUDevice, T>;            \
  REGISTER_KERNEL_BUILDER(                                       \
      Name("Example").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
      ExampleOp<GPUDevice, T>);
REGISTER_GPU(float);
REGISTER_GPU(int32);
#endif  // GOOGLE_CUDA

// kernel_example.cu.cc
#ifdef GOOGLE_CUDA
#define EIGEN_USE_GPU
#include "kernel_example.h"
#include "tensorflow/core/util/gpu_kernel_helper.h"

using namespace tensorflow;

using GPUDevice = Eigen::GpuDevice;

// Define the CUDA kernel.
template <typename T>
__global__ void ExampleCudaKernel(const int size, const T* in, T* out) {
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < size;
       i += blockDim.x * gridDim.x) {
    out[i] = 2 * __ldg(in + i);
  }
}

// Define the GPU implementation that launches the CUDA kernel.
template <typename T>
void ExampleFunctor<GPUDevice, T>::operator()(
    const GPUDevice& d, int size, const T* in, T* out) {
  // Launch the cuda kernel.
  //
  // See core/util/gpu_kernel_helper.h for example of computing
  // block count and thread_per_block count.
  int block_count = 1024;
  int thread_per_block = 20;
  ExampleCudaKernel<T>
      <<<block_count, thread_per_block, 0, d.stream()>>>(size, in, out);
}

// Explicitly instantiate functors for the types of OpKernels registered.
template struct ExampleFunctor<GPUDevice, float>;
template struct ExampleFunctor<GPUDevice, int32>;

#endif  // GOOGLE_CUDA

建構運算 (op) 程式庫

使用您的系統編譯器編譯運算 (op) (TensorFlow 二進位檔安裝)

您應該能夠使用系統上可用的 C++ 編譯器 (例如 g++ 或 clang) 編譯 zero_out.cc。二進位 PIP 套件會將您編譯運算 (op) 所需的標頭檔和程式庫安裝在系統特定的位置。但是，TensorFlow Python 程式庫提供 get_include 函式來取得標頭目錄，而 get_lib 目錄具有要連結的共用物件。以下是這些函式在 Ubuntu 機器上的輸出。

$ python
>>> import tensorflow as tf
>>> tf.sysconfig.get_include()
'/usr/local/lib/python3.6/site-packages/tensorflow/include'
>>> tf.sysconfig.get_lib()
'/usr/local/lib/python3.6/site-packages/tensorflow'

假設您已安裝 g++，以下是您可以用於將運算 (op) 編譯成動態程式庫的命令序列。

TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
TF_LFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))') )
g++ -std=c++14 -shared zero_out.cc -o zero_out.so -fPIC ${TF_CFLAGS[@]} ${TF_LFLAGS[@]} -O2

在 macOS 上，建構 .so 檔案時，需要額外的旗標「-undefined dynamic_lookup」。

關於 gcc 版本 >=5 的注意事項：gcc 自版本 5 起使用新的 C++ ABI。TensorFlow 2.8 和更早版本是使用 gcc4 建構的，後者使用較舊的 ABI。如果您使用的是這些版本的 TensorFlow，並且嘗試使用 gcc>=5 編譯您的運算 (op) 程式庫，請將 -D_GLIBCXX_USE_CXX11_ABI=0 新增至命令列，使程式庫與較舊的 ABI 相容。TensorFlow 2.9 以上版本套件預設與較新的 ABI 相容。

使用 bazel 編譯運算 (op) (TensorFlow 原始碼安裝)

如果您已安裝 TensorFlow 原始碼，則可以利用 TensorFlow 的建構系統來編譯您的運算 (op)。將包含以下 Bazel 建構規則的 BUILD 檔案放在 tensorflow/core/user_ops/ 目錄中。

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    name = "zero_out.so",
    srcs = ["zero_out.cc"],
)

執行以下命令以建構 zero_out.so。

$ bazel build --config opt //tensorflow/core/user_ops:zero_out.so

為了使用 CUDA 核心 (Kernel) 編譯 Example 運算，您需要使用 tf_custom_op_library 的 gpu_srcs 參數。將包含以下 Bazel 建構規則的 BUILD 檔案放在 tensorflow/core/user_ops/ 目錄中的新資料夾中 (例如，「example_gpu」)。

load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")

tf_custom_op_library(
    # kernel_example.cc  kernel_example.cu.cc  kernel_example.h
    name = "kernel_example.so",
    srcs = ["kernel_example.h", "kernel_example.cc"],
    gpu_srcs = ["kernel_example.cu.cc", "kernel_example.h"],
)

執行以下命令以建構 kernel_example.so。

$ bazel build --config opt //tensorflow/core/user_ops/example_gpu:kernel_example.so

在 Python 中使用運算 (op)

TensorFlow Python API 提供 tf.load_op_library 函式來載入動態程式庫，並向 TensorFlow 框架註冊運算 (op)。load_op_library 會傳回一個 Python 模組，其中包含運算 (op) 和核心 (kernel) 的 Python 封裝函式 (wrappers)。因此，一旦您建構了運算 (op)，就可以執行以下操作以從 Python 執行它

import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
print(zero_out_module.zero_out([[1, 2], [3, 4]]).numpy())

# Prints
array([[1, 0], [0, 0]], dtype=int32)

請記住，產生的函式將被賦予蛇紋命名法 (snake_case) 名稱 (以符合 PEP8)。因此，如果您的運算 (op) 在 C++ 檔案中命名為 ZeroOut，則 Python 函式將被呼叫為 zero_out。

若要使運算 (op) 可作為可從 Python 模組 import 的常規函式使用，在 Python 原始碼檔案中加入 load_op_library 呼叫可能會很有用，如下所示

import tensorflow as tf

zero_out_module = tf.load_op_library('./zero_out.so')
zero_out = zero_out_module.zero_out

驗證運算 (op) 是否運作

驗證您是否已成功實作運算 (op) 的一個好方法是為其編寫測試。建立檔案 zero_out_op_test.py，內容如下

import tensorflow as tf

class ZeroOutTest(tf.test.TestCase):
  def testZeroOut(self):
    zero_out_module = tf.load_op_library('./zero_out.so')
    with self.test_session():
      result = zero_out_module.zero_out([5, 4, 3, 2, 1])
      self.assertAllEqual(result.eval(), [5, 0, 0, 0, 0])

if __name__ == "__main__":
  tf.test.main()

然後執行您的測試 (假設您已安裝 tensorflow)

$ python zero_out_op_test.py

將進階功能建置到您的運算 (op) 中

現在您已了解如何建構基本 (且在某種程度上受限制) 的運算 (op) 和實作，我們將研究您通常需要建置到運算 (op) 中的一些更複雜的事項。這包括

• 條件檢查和驗證
• 運算 (op) 註冊
  • 屬性 (Attrs)
  • 屬性 (Attr) 類型
  • 多型性
  • 輸入和輸出
  • 回溯相容性
• GPU 支援
  • 編譯 GPU 裝置的核心 (kernel)
• 在 Python 中實作梯度
• C++ 中的形狀函式

條件檢查和驗證

上面的範例假設運算 (op) 適用於任何形狀的張量。如果它僅適用於向量呢？這表示需要將檢查新增至上述 OpKernel 實作。

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);

    OP_REQUIRES(context, TensorShapeUtils::IsVector(input_tensor.shape()),
                errors::InvalidArgument("ZeroOut expects a 1-D vector."));
    // ...
  }

這會斷言輸入是向量，如果不是，則傳回並設定 InvalidArgument 狀態。OP_REQUIRES 巨集接受三個引數

• context，可以是 OpKernelContext 或 OpKernelConstruction 指標 (請參閱 tensorflow/core/framework/op_kernel.h)，用於其 SetStatus() 方法。
• 條件。例如，在 tensorflow/core/framework/tensor_shape.h 中有用於驗證張量形狀的函式
• 錯誤本身，由 Status 物件表示，請參閱 tensorflow/core/platform/status.h。Status 具有類型 (通常為 InvalidArgument，但請參閱類型清單) 和訊息。用於建構錯誤的函式可以在 tensorflow/core/platform/errors.h 中找到。

或者，如果您想測試從某個函式傳回的 Status 物件是否為錯誤，如果是，則傳回它，請使用 OP_REQUIRES_OK。這兩個巨集都會在發生錯誤時從函式傳回。

運算 (op) 註冊

屬性 (Attrs)

運算 (ops) 可以具有屬性 (attrs)，其值會在將運算 (op) 新增至圖時設定。這些屬性 (attrs) 用於設定運算 (op)，並且可以在核心 (kernel) 實作中以及運算 (op) 註冊中輸入和輸出的類型中存取它們的值。在可能的情況下，最好使用輸入而不是屬性 (attr)，因為輸入更具彈性。這是因為屬性 (attrs) 是常數，必須在圖建構時定義。相反地，輸入是張量 (Tensors)，其值可以是動態的；也就是說，輸入可以每步變更、使用饋送等方式設定。屬性 (Attrs) 用於無法使用輸入執行的操作：任何影響簽名 (輸入或輸出的數量或類型) 或無法逐步變更的設定。

您可以在註冊運算 (op) 時定義屬性 (attr)，方法是使用 Attr 方法指定其名稱和類型，該方法預期規格的形式為

<name>: <attr-type-expr>

其中 <name> 以字母開頭，並且可以由字母數字字元和底線組成，而 <attr-type-expr> 是以下描述的形式的類型運算式。

例如，如果您希望 ZeroOut 運算 (op) 保留使用者指定的索引，而不是僅保留第 0 個元素，則可以像這樣註冊運算 (op)

REGISTER_OP("ZeroOut")
    .Attr("preserve_index: int")
    .Input("to_zero: int32")
    .Output("zeroed: int32");

(請注意，屬性類型集合與用於輸入和輸出的 tf.DType 不同。)

然後，您的核心 (kernel) 可以透過 context 參數在其建構函式中存取此屬性 (attr)

class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {
    // Get the index of the value to preserve
    OP_REQUIRES_OK(context,
                   context->GetAttr("preserve_index", &preserve_index_));
    // Check that preserve_index is positive
    OP_REQUIRES(context, preserve_index_ >= 0,
                errors::InvalidArgument("Need preserve_index >= 0, got ",
                                        preserve_index_));
  }
  void Compute(OpKernelContext* context) override {
    // ...
  }
 private:
  int preserve_index_;
};

然後可以在 Compute 方法中使用它

  void Compute(OpKernelContext* context) override {
    // ...

    // We're using saved attr to validate potentially dynamic input
    // So we check that preserve_index is in range
    OP_REQUIRES(context, preserve_index_ < input.dimension(0),
                errors::InvalidArgument("preserve_index out of range"));

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the requested input value
    output_flat(preserve_index_) = input(preserve_index_);
  }

屬性 (Attr) 類型

屬性 (attr) 中支援以下類型

• string：任何位元組序列 (不一定為 UTF8)。
• int：帶正負號的整數。
• float：浮點數。
• bool：True 或 false。
• type：DataType 的 (非 ref) 值之一。
• shape：TensorShapeProto。
• list(<type>)：<type> 的清單，其中 <type> 是上述類型之一。請注意，list(list(<type>)) 無效。

另請參閱：op_def_builder.cc:FinalizeAttr 以取得明確的清單。

預設值和限制

屬性 (Attrs) 可以具有預設值，並且某些類型的屬性 (attrs) 可以具有限制。若要定義具有限制的屬性 (attr)，您可以使用以下 <attr-type-expr>。

{'<string1>', '<string2>'}：值必須是字串，其值為 <string1> 或 <string2>。當您使用此語法時，類型名稱 string 是隱含的。這會模擬列舉

REGISTER_OP("EnumExample")
    .Attr("e: {'apple', 'orange'}");

{<type1>, <type2>}：值類型為 type，並且必須是 <type1> 或 <type2> 之一，其中 <type1> 和 <type2> 是支援的 tf.DType。您無需指定屬性 (attr) 的類型為 type。當您在 {...} 中具有類型清單時，這是隱含的。例如，在此情況下，屬性 (attr) t 是一種必須是 int32、float 或 bool 的類型

REGISTER_OP("RestrictedTypeExample")
    .Attr("t: {int32, float, bool}");

常見的類型限制有捷徑

• numbertype：類型 type 限制為數值 (非字串和非布林值) 類型。
• realnumbertype：類似於 numbertype，但不包含複數類型。
• quantizedtype：類似於 numbertype，但僅包含量化數字類型。

這些允許的特定類型清單由 tensorflow/core/framework/types.h 中的函式 (例如 NumberTypes()) 定義。在此範例中，屬性 (attr) t 必須是數值類型之一

REGISTER_OP("NumberType")
    .Attr("t: numbertype");

對於此運算 (op)

tf.number_type(t=tf.int32)  # Valid
tf.number_type(t=tf.bool)   # Invalid

清單可以與其他清單和單一類型組合。以下運算 (op) 允許屬性 (attr) t 為任何數值類型或布林值類型

REGISTER_OP("NumberOrBooleanType")
    .Attr("t: {numbertype, bool}");

對於此運算 (op)

tf.number_or_boolean_type(t=tf.int32)  # Valid
tf.number_or_boolean_type(t=tf.bool)   # Valid
tf.number_or_boolean_type(t=tf.string) # Invalid

int >= <n>：值必須是整數，其值大於或等於 <n>，其中 <n> 是自然數。例如，以下運算 (op) 註冊指定屬性 (attr) a 的值必須至少為 2

REGISTER_OP("MinIntExample")
    .Attr("a: int >= 2");

list(<type>) >= <n>：類型為 <type> 的清單，其長度大於或等於 <n>。例如，以下運算 (op) 註冊指定屬性 (attr) a 是類型清單 (int32 或 float)，並且必須至少有 3 個

REGISTER_OP("TypeListExample")
    .Attr("a: list({int32, float}) >= 3");

若要為屬性 (attr) 設定預設值 (使其在產生的程式碼中成為選用項目)，請在結尾新增 = <default>，如下所示

REGISTER_OP("AttrDefaultExample")
    .Attr("i: int = 0");

此外，可以同時指定限制和預設值

REGISTER_OP("AttrConstraintAndDefaultExample")
    .Attr("i: int >= 1 = 1");

預設值的支援語法將用於結果 GraphDef 定義的 proto 表示法中。

以下是如何為所有類型指定預設值的範例

REGISTER_OP("AttrDefaultExampleForAllTypes")
   .Attr("s: string = 'foo'")
   .Attr("i: int = 0")
   .Attr("f: float = 1.0")
   .Attr("b: bool = true")
   .Attr("ty: type = DT_INT32")
   .Attr("sh: shape = { dim { size: 1 } dim { size: 2 } }")
   .Attr("te: tensor = { dtype: DT_INT32 int_val: 5 }")
   .Attr("l_empty: list(int) = []")
   .Attr("l_int: list(int) = [2, 3, 5, 7]");

特別注意，type 類型的值使用 tf.DType。

多型性

類型多型性

對於可以接受不同類型作為輸入或產生不同輸出類型的運算 (ops)，您可以在運算 (op) 註冊中的輸入或輸出類型中指定屬性 (attr)。通常，您會為每個支援的類型註冊一個 OpKernel。

例如，如果您希望 ZeroOut 運算 (op) 除了 int32 之外，還能處理 float，則您的運算 (op) 註冊可能看起來像這樣

REGISTER_OP("ZeroOut")
    .Attr("T: {float, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

您的運算 (op) 註冊現在指定輸入的類型必須是 float 或 int32，並且其輸出將是相同的類型，因為兩者都具有類型 T。

命名

輸入、輸出和屬性 (attrs) 通常應給予蛇紋命名法 (snake_case) 名稱。唯一的例外是屬性 (attrs)，它們用作輸入的類型或輸出類型。當運算 (op) 新增至圖時，可以推斷這些屬性 (attrs)，因此不會出現在運算 (op) 的函式中。例如，ZeroOut 的最後一個定義將產生一個類似於以下的 Python 函式

def zero_out(to_zero, name=None):
  """...
  Args:
    to_zero: A `Tensor`. Must be one of the following types:
        `float32`, `int32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `to_zero`.
  """

如果傳遞 to_zero 一個 int32 張量，則 T 會自動設定為 int32 (實際上是 DT_INT32)。這些推斷出的屬性會以大寫或駝峰式命名。

將此與具有決定輸出類型之類型屬性的運算進行比較

REGISTER_OP("StringToNumber")
    .Input("string_tensor: string")
    .Output("output: out_type")
    .Attr("out_type: {float, int32} = DT_FLOAT");
    .Doc(R"doc(
Converts each string in the input Tensor to the specified numeric type.
)doc");

在這種情況下，使用者必須指定輸出類型，如產生的 Python 所示

def string_to_number(string_tensor, out_type=None, name=None):
  """Converts each string in the input Tensor to the specified numeric type.

  Args:
    string_tensor: A `Tensor` of type `string`.
    out_type: An optional `tf.DType` from: `tf.float32, tf.int32`.
      Defaults to `tf.float32`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor` of type `out_type`.
  """

類型多型範例

#include "tensorflow/core/framework/op_kernel.h"

class ZeroOutInt32Op : public OpKernel {
  // as before
};

class ZeroOutFloatOp : public OpKernel {
 public:
  explicit ZeroOutFloatOp(OpKernelConstruction* context)
      : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<float>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<float>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutInt32Op);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutFloatOp);

為了保持向後相容性，當新增屬性到現有運算時，您應指定預設值

REGISTER_OP("ZeroOut")
  .Attr("T: {float, int32} = DT_INT32")
  .Input("to_zero: T")
  .Output("zeroed: T")

假設您想要新增更多類型，例如 double

REGISTER_OP("ZeroOut")
    .Attr("T: {float, double, int32}")
    .Input("to_zero: T")
    .Output("zeroed: T");

您可以改用 C++ 範本，而不是像上面那樣編寫另一個具有冗餘程式碼的 OpKernel。每個多載版本仍然會有一個核心註冊 (REGISTER_KERNEL_BUILDER 呼叫)。

template <typename T>
class ZeroOutOp : public OpKernel {
 public:
  explicit ZeroOutOp(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    // Grab the input tensor
    const Tensor& input_tensor = context->input(0);
    auto input = input_tensor.flat<T>();

    // Create an output tensor
    Tensor* output = NULL;
    OP_REQUIRES_OK(context,
                   context->allocate_output(0, input_tensor.shape(), &output));
    auto output_flat = output->template flat<T>();

    // Set all the elements of the output tensor to 0
    const int N = input.size();
    for (int i = 0; i < N; i++) {
      output_flat(i) = 0;
    }

    // Preserve the first input value
    if (N > 0) output_flat(0) = input(0);
  }
};

// Note that TypeConstraint<int32>("T") means that attr "T" (defined
// in the op registration above) must be "int32" to use this template
// instantiation.
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<int32>("T"),
    ZeroOutOp<int32>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<float>("T"),
    ZeroOutOp<float>);
REGISTER_KERNEL_BUILDER(
    Name("ZeroOut")
    .Device(DEVICE_CPU)
    .TypeConstraint<double>("T"),
    ZeroOutOp<double>);

如果您有多個多載版本，您可以將註冊放在巨集中。

#include "tensorflow/core/framework/op_kernel.h"

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

REGISTER_KERNEL(int32);
REGISTER_KERNEL(float);
REGISTER_KERNEL(double);

#undef REGISTER_KERNEL

根據您為核心註冊的類型列表，您或許可以使用 tensorflow/core/framework/register_types.h 提供的巨集。

#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"

REGISTER_OP("ZeroOut")
    .Attr("T: realnumbertype")
    .Input("to_zero: T")
    .Output("zeroed: T");

template <typename T>
class ZeroOutOp : public OpKernel { ... };

#define REGISTER_KERNEL(type)                                       \
  REGISTER_KERNEL_BUILDER(                                          \
      Name("ZeroOut").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
      ZeroOutOp<type>)

TF_CALL_REAL_NUMBER_TYPES(REGISTER_KERNEL);

#undef REGISTER_KERNEL

列出輸入和輸出

除了能夠接受或產生不同類型之外，運算還可以消耗或產生可變數量的張量。

在下一個範例中，屬性 T 保留了類型列表，並用作輸入 in 和輸出 out 的類型。輸入和輸出是該類型的張量列表 (而且輸出中張量的數量和類型與輸入相同，因為兩者都具有類型 T)。

REGISTER_OP("PolymorphicListExample")
    .Attr("T: list(type)")
    .Input("in: T")
    .Output("out: T");

您也可以限制列表中可以指定的類型。在下一個範例中，輸入是 float 和 double 張量的列表。例如，此運算接受輸入類型 (float, double, float)，在這種情況下，輸出類型也將是 (float, double, float)。

REGISTER_OP("ListTypeRestrictionExample")
    .Attr("T: list({float, double})")
    .Input("in: T")
    .Output("out: T");

如果您希望列表中的所有張量都屬於同一類型，您可以執行類似以下的操作

REGISTER_OP("IntListInputExample")
    .Attr("N: int")
    .Input("in: N * int32")
    .Output("out: int32");

這會接受 int32 張量的列表，並使用 int 屬性 N 來指定列表的長度。

這也可以製成類型多型。在下一個範例中，輸入是相同 (但未指定) 類型 ("T") 的張量列表 (長度為 "N")，而輸出是單個相符類型的張量

REGISTER_OP("SameListInputExample")
    .Attr("N: int")
    .Attr("T: type")
    .Input("in: N * T")
    .Output("out: T");

預設情況下，張量列表的最小長度為 1。您可以使用對應屬性上的 ">=" 限制來變更該預設值。在下一個範例中，輸入是至少 2 個 int32 張量的列表

REGISTER_OP("MinLengthIntListExample")
    .Attr("N: int >= 2")
    .Input("in: N * int32")
    .Output("out: int32");

相同的語法適用於 "list(type)" 屬性

REGISTER_OP("MinimumLengthPolymorphicListExample")
    .Attr("T: list(type) >= 3")
    .Input("in: T")
    .Output("out: T");

輸入和輸出

總結以上內容，運算註冊可以有多個輸入和輸出

REGISTER_OP("MultipleInsAndOuts")
    .Input("y: int32")
    .Input("z: float")
    .Output("a: string")
    .Output("b: int32");

每個輸入或輸出規格的形式為

<name>: <io-type-expr>

其中 <name> 以字母開頭，並且可以由字母數字字元和底線組成。<io-type-expr> 是以下類型運算式之一

• <type>，其中 <type> 是支援的輸入類型 (例如 float、int32、string)。這指定了給定類型的單個張量。

  請參閱 tf.DType。

  REGISTER_OP("BuiltInTypesExample")
      .Input("integers: int32")
      .Input("complex_numbers: complex64");
  

• <attr-type>，其中 <attr-type> 是類型為 type 或 list(type) (可能具有類型限制) 的 Attr 的名稱。此語法允許多型運算。

  REGISTER_OP("PolymorphicSingleInput")
      .Attr("T: type")
      .Input("in: T");
  
  REGISTER_OP("RestrictedPolymorphicSingleInput")
      .Attr("T: {int32, int64}")
      .Input("in: T");
  

  參考類型為 list(type) 的屬性可讓您接受張量序列。

  REGISTER_OP("ArbitraryTensorSequenceExample")
      .Attr("T: list(type)")
      .Input("in: T")
      .Output("out: T");
  
  REGISTER_OP("RestrictedTensorSequenceExample")
      .Attr("T: list({int32, int64})")
      .Input("in: T")
      .Output("out: T");
  

  請注意，輸出 out 中張量的數量和類型與輸入 in 相同，因為兩者都屬於類型 T。

• 對於具有相同類型的張量序列：<number> * <type>，其中 <number> 是類型為 int 的 Attr 的名稱。<type> 可以是 tf.DType，也可以是類型為 type 的屬性名稱。作為第一個範例，此運算接受 int32 張量的列表

  REGISTER_OP("Int32SequenceExample")
      .Attr("NumTensors: int")
      .Input("in: NumTensors * int32")
  

  而此運算接受任何類型的張量列表，只要它們都相同

  REGISTER_OP("SameTypeSequenceExample")
      .Attr("NumTensors: int")
      .Attr("T: type")
      .Input("in: NumTensors * T")
  

• 對於張量的參考：Ref(<type>)，其中 <type> 是先前的類型之一。

輸入類型中使用的任何屬性都會被推斷出來。依照慣例，這些推斷出的屬性會使用大寫名稱 (如 T 或 N)。否則，輸入、輸出和屬性的名稱會像函式參數一樣 (例如 num_outputs)。如需更多詳細資訊，請參閱先前的命名章節。

如需更多詳細資訊，請參閱 tensorflow/core/framework/op_def_builder.h。

回溯相容性

假設您編寫了一個出色的自訂運算並與他人分享，因此您有滿意的客戶使用您的運算。但是，您想以某種方式變更運算。

一般而言，對現有、已簽入規格的變更必須向後相容：變更運算的規格不得破壞先前從較舊規格建構的序列化 GraphDef 通訊協定緩衝區。GraphDef 相容性的詳細資訊在此處說明。

有幾種方法可以保持向後相容性。

1. 新增至運算的任何新屬性都必須定義預設值，並且使用該預設值，運算必須具有原始行為。若要將運算從非多型變更為多型，您必須為新的類型屬性提供預設值，以預設方式保留原始簽名。例如，如果您的運算是

   REGISTER_OP("MyGeneralUnaryOp")
       .Input("in: float")
       .Output("out: float");
   

   您可以使用以下方式以向後相容的方式使其成為多型

   REGISTER_OP("MyGeneralUnaryOp")
       .Input("in: T")
       .Output("out: T")
       .Attr("T: numerictype = DT_FLOAT");
   

2. 您可以安全地使屬性上的限制變得較不嚴格。例如，您可以從 {int32, int64} 變更為 {int32, int64, float} 或 type。或者您可以從 {"apple", "orange"} 變更為 {"apple", "banana", "orange"} 或 string。

3. 您可以將單個輸入/輸出變更為列表輸入/輸出，只要列表類型的預設值與舊簽名相符即可。

4. 您可以新增新的列表輸入/輸出，如果預設為空。

5. 命名空間您建立的任何新運算，方法是在運算名稱前面加上專案獨有的內容。這可避免您的運算與未來 TensorFlow 版本中可能包含的任何運算發生衝突。

6. 提前規劃！嘗試預測運算未來的用途。某些簽名變更無法以相容的方式完成 (例如，將相同類型的列表變更為不同類型的列表)。

安全和不安全變更的完整列表可以在 tensorflow/core/framework/op_compatibility_test.cc 中找到。如果您無法使對運算的變更向後相容，請使用具有新語意的新名稱建立新運算。

另請注意，雖然這些變更可以維持 GraphDef 相容性，但產生的 Python 程式碼可能會以與舊呼叫者不相容的方式變更。可以透過手寫 Python 包裝函式中的仔細變更，以及保留舊簽名 (除了可能在結尾新增新的選用引數) 來保持 Python API 相容性。通常，不相容的變更只能在 TensorFlow 變更主要版本時進行，並且必須符合GraphDef 版本語意。

GPU 支援

您可以實作不同的 OpKernel，並為 CPU 和 GPU 各註冊一個，就像您可以為不同類型註冊核心一樣。在 tensorflow/core/kernels/ 中有幾個具有 GPU 支援的核心範例。請注意，某些核心在 .cc 檔案中具有 CPU 版本，在以 _gpu.cu.cc 結尾的檔案中具有 GPU 版本，以及在 .h 檔案中具有一些共用程式碼。

例如，tf.pad 在 tensorflow/core/kernels/pad_op.cc 中具有除 GPU 核心之外的所有內容。GPU 核心位於 tensorflow/core/kernels/pad_op_gpu.cu.cc 中，而共用程式碼是一個範本類別，定義於 tensorflow/core/kernels/pad_op.h 中。我們以這種方式組織程式碼的原因有兩個：它可讓您在 CPU 和 GPU 實作之間共用通用程式碼，並且它將 GPU 實作放在單獨的檔案中，以便只能由 GPU 編譯器編譯。

需要注意的一件事是，即使使用 pad 的 GPU 核心版本，它仍然需要在 CPU 記憶體中使用 "paddings" 輸入。若要標記輸入或輸出保留在 CPU 上，請將 HostMemory() 呼叫新增至核心註冊，例如

#define REGISTER_GPU_KERNEL(T)                         \
  REGISTER_KERNEL_BUILDER(Name("Pad")                  \
                              .Device(DEVICE_GPU)      \
                              .TypeConstraint<T>("T")  \
                              .HostMemory("paddings"), \
                          PadOp<GPUDevice, T>)

編譯 GPU 裝置的核心 (kernel)

查看 cuda_op_kernel.cu.cc 以取得使用 CUDA 核心實作運算的範例。tf_custom_op_library 接受 gpu_srcs 引數，可以在其中指定包含 CUDA 核心 (*.cu.cc 檔案) 的來源檔案列表。若要與 TensorFlow 的二進位安裝搭配使用，CUDA 核心必須使用 NVIDIA 的 nvcc 編譯器進行編譯。以下是您可以用來將 cuda_op_kernel.cu.cc 和 cuda_op_kernel.cc 編譯成單個動態可載入程式庫的命令序列

nvcc -std=c++14 -c -o cuda_op_kernel.cu.o cuda_op_kernel.cu.cc \
  ${TF_CFLAGS[@]} -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC

g++ -std=c++14 -shared -o cuda_op_kernel.so cuda_op_kernel.cc \
  cuda_op_kernel.cu.o ${TF_CFLAGS[@]} -fPIC -lcudart ${TF_LFLAGS[@]}

上面產生的 cuda_op_kernel.so 可以像往常一樣在 Python 中載入，使用 tf.load_op_library 函式。

請注意，如果您的 CUDA 程式庫未安裝在 /usr/local/lib64 中，您需要在第二個 (g++) 命令中明確指定路徑。例如，如果您的 CUDA 安裝在 /usr/local/cuda-8.0 中，請新增 -L /usr/local/cuda-8.0/lib64/。

在 Python 中實作梯度

給定運算的圖形，TensorFlow 使用自動微分 (反向傳播) 來新增代表相對於現有運算之梯度的運算。若要使自動微分適用於新運算，您必須註冊梯度函式，該函式計算相對於運算輸入的梯度，並給定相對於運算輸出的梯度。

在數學上，如果運算計算 \(y = f(x)\)，則註冊的梯度運算會透過鏈式法則將損失 \(L\) 相對於 \(y\) 的梯度 \(\partial L/ \partial y\) 轉換為相對於 \(x\) 的梯度 \(\partial L/ \partial x\)

\[\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial f}{\partial x}.\]

在 ZeroOut 的情況下，只有輸入中的一個條目會影響輸出，因此相對於輸入的梯度是稀疏的「單熱」張量。這表示如下

from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import sparse_ops

@ops.RegisterGradient("ZeroOut")
def _zero_out_grad(op, grad):
  """The gradients for `zero_out`.

  Args:
    op: The `zero_out` `Operation` that we are differentiating, which we can use
      to find the inputs and outputs of the original op.
    grad: Gradient with respect to the output of the `zero_out` op.

  Returns:
    Gradients with respect to the input of `zero_out`.
  """
  to_zero = op.inputs[0]
  shape = array_ops.shape(to_zero)
  index = array_ops.zeros_like(shape)
  first_grad = array_ops.reshape(grad, [-1])[0]
  to_zero_grad = sparse_ops.sparse_to_dense([index], shape, first_grad, 0)
  return [to_zero_grad]  # List of one Tensor, since we have one input

有關使用 tf.RegisterGradient 註冊梯度函式的詳細資訊

• 對於具有一個輸出的運算，梯度函式將採用 tf.Operation、op 和 tf.Tensor grad，並從張量 op.inputs[i]、op.outputs[i] 和 grad 建立新的運算。有關任何屬性的資訊可以透過 tf.Operation.get_attr 找到。

• 如果運算有多個輸出，則梯度函式將採用 op 和 grads，其中 grads 是相對於每個輸出的梯度列表。梯度函式的結果必須是代表相對於每個輸入之梯度的 Tensor 物件列表。

• 如果某些輸入沒有明確定義的梯度 (例如用作索引的整數輸入)，則對應的傳回梯度應為 None。例如，對於採用浮點張量 x 和整數索引 i 的運算，梯度函式將 return [x_grad, None]。

• 如果運算根本沒有有意義的梯度，您通常不必註冊任何梯度，而且只要永遠不需要運算的梯度，您就可以正常運作。在某些情況下，運算沒有明確定義的梯度，但可能參與梯度的計算。在這裡，您可以使用 ops.NotDifferentiable 來自動向後傳播零。

請注意，在呼叫梯度函式時，只有運算的資料流程圖可用，而不是張量資料本身。因此，所有計算都必須使用其他 tensorflow 運算來執行，以便在圖形執行時執行。

在為運算類型註冊自訂梯度時新增類型提示，使程式碼更具可讀性、更易於偵錯、更易於維護，並透過資料驗證更穩健。例如，當在函式中採用 op 作為參數時，請指定梯度函式將採用 tf.Operation 作為其參數類型。

C++ 中的形狀函式

TensorFlow API 具有稱為「形狀推斷」的功能，可提供有關張量形狀的資訊，而無需執行圖形。形狀推斷由「形狀函式」支援，這些函式在 C++ REGISTER_OP 宣告中為每個運算類型註冊，並執行兩個角色：在圖形建構期間斷言輸入的形狀相容，以及指定輸出的形狀。

形狀函式定義為對 shape_inference::InferenceContext 類別的運算。例如，在 ZeroOut 的形狀函式中

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

c->set_output(0, c->input(0)); 宣告第一個輸出的形狀應設定為第一個輸入的形狀。如果輸出是依其索引選取的 (如上述範例所示)，則 set_output 的第二個參數應為 ShapeHandle 物件。您可以透過其預設建構函式建立空的 ShapeHandle 物件。索引為 idx 的輸入的 ShapeHandle 物件可以透過 c->input(idx) 取得。

有許多適用於許多運算的常見形狀函式，例如 shape_inference::UnchangedShape，可以在 common_shape_fns.h 中找到，並按如下方式使用

REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);

形狀函式也可以限制輸入的形狀。對於具有向量形狀限制的 ZeroOut 版本，形狀函式將如下所示

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &input));
      c->set_output(0, input);
      return Status::OK();
    });

WithRank 呼叫驗證輸入形狀 c->input(0) 是否具有恰好一個維度的形狀 (或者，如果輸入形狀未知，則輸出形狀將是一個具有一個未知維度的向量)。

如果您的運算是具有多個輸入的多型，您可以使用 InferenceContext 的成員來判斷要檢查的形狀數量，並使用 Merge 來驗證形狀是否都相容 (或者，存取指示長度的屬性，使用 InferenceContext::GetAttr，它提供對運算屬性的存取權)。

    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      ::tensorflow::shape_inference::ShapeHandle input;
      ::tensorflow::shape_inference::ShapeHandle output;
      for (size_t i = 0; i < c->num_inputs(); ++i) {
        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 2, &input));
        TF_RETURN_IF_ERROR(c->Merge(output, input, &output));
      }
      c->set_output(0, output);
      return Status::OK();
    });

由於形狀推斷是選用功能，而且張量的形狀可能會動態變化，因此形狀函式必須能穩健地處理任何輸入的不完整形狀資訊。InferenceContext 中的 Merge 方法允許呼叫者斷言兩個形狀相同，即使它們中的一個或兩個都沒有完整資訊。形狀函式是為所有核心 TensorFlow 運算定義的，並提供許多不同的使用範例。

InferenceContext 類別具有許多可用於定義形狀函式操作的函式。例如，您可以使用 InferenceContext::Dim 和 InferenceContext::WithValue 驗證特定維度是否具有非常特定的值；您可以使用 InferenceContext::Add 和 InferenceContext::Multiply 指定輸出維度是兩個輸入維度的總和/乘積。請參閱 InferenceContext 類別，以瞭解您可以指定的所有各種形狀操作。以下範例將第一個輸出的形狀設定為 (n, 3)，其中第一個輸入的形狀為 (n, ...)

.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
    c->set_output(0, c->Matrix(c->Dim(c->input(0), 0), 3));
    return Status::OK();
});

如果您有複雜的形狀函式，您應考慮新增測試以驗證各種輸入形狀組合是否產生預期的輸出形狀組合。您可以在我們的一些核心運算測試中看到如何編寫這些測試的範例。(INFER_OK 和 INFER_ERROR 的語法有點隱晦，但請盡量簡潔地表示測試中的輸入和輸出形狀規格。目前，請參閱這些測試中的周圍註解，以瞭解形狀字串規格的含義)。

為您的自訂運算建立 pip 套件

若要為您的運算建立 pip 套件，請參閱 tensorflow/custom-op 範例。本指南說明如何從 TensorFlow pip 套件建置自訂運算，而不是從原始碼建置 TensorFlow。