No history yet

Op Registration API

Defining an Op's Interface

When you create a custom operation in TensorFlow, the first step is to declare its interface. This is the contract that tells the TensorFlow runtime what your op is called, what data it expects, and what it produces. This entire definition happens in C++ and is handled by the REGISTER_OP macro.

Think of it as a function prototype for the computation graph. It defines the op's signature, which the graph compiler uses to validate connections and prepare for execution, all without knowing the specific implementation details (the kernel).

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

REGISTER_OP("MyCustomOp");

This is the simplest possible registration. It just tells TensorFlow that an operation named "MyCustomOp" exists. It doesn't take any inputs or produce any outputs yet.

To make it useful, we chain calls to the REGISTER_OP macro to specify its inputs and outputs. The .Input() and .Output() methods take a string that defines the argument name and its data type, separated by a colon.

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

REGISTER_OP("AddOne")
    .Input("input_tensor: int32")
    .Output("output_tensor: int32");

Handling Different Data Types

Hardcoding a single data type like int32 is restrictive. Often, you want an operation to work on float, double, and int tensors. This is achieved through type polymorphism using attributes, defined with the .Attr() method. An attribute is a parameter that is fixed when the op is added to the graph.

A common pattern is to define an attribute of type type. This creates a placeholder for a data type that can be referenced in the input and output definitions. By convention, this attribute is often named T.

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

REGISTER_OP("AddOnePolymorphic")
    .Attr("T: type")
    .Input("input_tensor: T")
    .Output("output_tensor: T");

Here, T is a placeholder for the type. When someone uses this op in Python, TensorFlow will infer the type from the input tensor (tf.float32, tf.int64, etc.) and substitute it for T. You can also add constraints to limit which types are allowed and set default values for attributes.

REGISTER_OP("AddOneConstrained")
    // T can only be one of these numeric types
    .Attr("T: {float, double, int32, int64}")
    // An integer attribute with a default value
    .Attr("increment_by: int = 1")
    .Input("input_tensor: T")
    .Output("output_tensor: T");

Shape Inference: The Graph's Blueprint

One of the most critical parts of an op's interface is its shape function. TensorFlow builds a static graph, and to do so, it must be able to infer the shape of each tensor without actually running any computations. This allows it to detect shape-related errors early and perform optimizations. You define this logic using the .SetShapeFn() method.

The shape function receives a pointer to a shape_inference::InferenceContext, which we'll call c. This context object provides access to input shapes and methods for setting output shapes. A shape itself is represented by a ShapeHandle object.

For our simple "AddOne" op, the output shape is always identical to the input shape. The shape function is straightforward: it gets the shape of the first input and sets it as the shape for the first output.

using namespace tensorflow;

REGISTER_OP("AddOne")
    .Input("input_tensor: int32")
    .Output("output_tensor: int32")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      // Get the shape of the first input (index 0)
      shape_inference::ShapeHandle input_shape = c->input(0);
      
      // Set the shape of the first output (index 0)
      c->set_output(0, input_shape);
      
      return Status::OK();
    });

You can also perform validation inside the shape function. For example, you might require that an input tensor must be a vector (rank 1) or a matrix (rank 2). The InferenceContext provides methods for these checks.

// ... continuing from the previous example
.SetShapeFn([](shape_inference::InferenceContext* c) {
  shape_inference::ShapeHandle input_shape;
  
  // TF_RETURN_IF_ERROR is a macro that checks the Status
  // and returns immediately if it's an error.
  // Here we verify the input has a rank of 2.
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &input_shape));
  
  c->set_output(0, input_shape);
  return Status::OK();
});

By defining the op's name, its inputs and outputs, its attributes, and its shape function, you have created a complete interface. TensorFlow now knows how your op fits into the larger computation graph, even before you've written the first line of its hardware-specific kernel implementation.