Tensorflow migration problem regardless defining the operations tensorflow.python.framework.ops

Disclaimer: Unless no other choice, I would strongly recommend sticking to TensorFlow 1.x. Migrating code from TF 1.x to 2.x can be incredibly time consuming.

Registering a shape is done in c++ using SetShapeFn and not in python since TF 1.0. However, the private python API was kept in TF 1.x (I presume for backward compatibility reasons), but was completely removed in TF 2.0.

The Create an Op guide is really useful to migrate the code in that case, and I would really recommend reading it.

First of all, why is it needed to register a shape? It's needed for shape inference, a feature that allow TensorFlow to know the shape of the inputs and outputs in the computation graph without running actual code. Shape Inference allows, for example, error handling when trying to use an op on Tensors that don't have a compatible shape.

In your specific case, you need to convert the python code that uses ops.RegisterShape to c++ code using SetShapeFn. Thankfully, the github repository you are working with provides comments that are going to be helpful.

Let's start with the approx_match function. The python code is the following:

def approx_match(xyz1,xyz2):
    '''
input:
    xyz1 : batch_size * #dataset_points * 3
    xyz2 : batch_size * #query_points * 3
returns:
    match : batch_size * #query_points * #dataset_points
    '''
    return approxmatch_module.approx_match(xyz1,xyz2)
@ops.RegisterShape('ApproxMatch')
def _approx_match_shape(op):
    shape1=op.inputs[0].get_shape().with_rank(3)
    shape2=op.inputs[1].get_shape().with_rank(3)
    return [tf.TensorShape([shape1.dims[0],shape2.dims[1],shape1.dims[1]])]

Reading the code and the comments, we understand the following:

• There is 2 inputs xyz1, and xyz2
• xyz1 has the shape (batch_size, dataset_points, 3)
• xyz2 has the shape (batch_size, query_points, 3)
• There is 1 output: match
• match has the shape (batch_size, query_points, dataset_points)

This would translate to the following C++ code:

#include "tensorflow/core/framework/shape_inference.h"
using namespace tensorflow;
REGISTER_OP("ApproxMatch")
    .Input("xyz1: float32")
    .Input("xyz2: float32")
    .Output("match: float32")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
        shape_inference::ShapeHandle xyz1_shape = c->input(0);
        shape_inference::ShapeHandle xyz2_shape = c->input(1);
        // batch_size is the first dimension 
        shape_inference::DimensionHandle batch_size = c->Dim(xyz1_shape, 0);
        // dataset_points points is the 2nd dimension of the first input
        shape_inference::DimensionHandle dataset_points = c->Dim(xyz1_shape, 1);
        // query_points points is the 2nd dimension of the second input
        shape_inference::DimensionHandle query_points = c->Dim(xyz2_shape, 1);
        // Creating the new shape (batch_size, query_points, dataset_points)
        // and setting it to the output
        c->set_output(0, c->MakeShape({batch_size, query_points, dataset_points})); 
        // Returning a status telling that everything went well
        return Status::OK();    
    });

Warning: This code does not contain any error handling (for example, checking that the first dimensions of the 2 inputs are the same, or that the last dimension of both inputs is 3). I let that as an exercise to the reader, you could look at the aforementioned guide or directly to the source code of some ops to have an idea on how to do error handling, with, for example, the macro TF_RETURN_IF_ERROR.

The same steps can be applied to the match_cost function, which could look like that:

REGISTER_OP("MatchCost")
    .Input("xyz1: float32")
    .Input("xyz2: float32")
    .Input("match: float32")
    .Output("cost: float32")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
        shape_inference::DimensionHandle batch_size = c->Dim(c->input(0), 0);
        c->set_output(0, c->Vector(batch_size));
        return Status::OK();        
    });

You need then to recompile the so library with the makefile included in the project. You might have to change some flags, for example, TF 2.8 uses the c++ standard c++14 instead of c++11, so the flag -std=c++14 is needed. Once the library is compiled, you can test importing it in python:

>>> import tensorflow as tf
>>> approxmatch_module = tf.load_op_library('./tf_approxmatch_so.so')
>>> a = tf.random.uniform((10,20,3))
>>> b = tf.random.uniform((10,50,3))
>>> c = approxmatch_module.approx_match(a,b)
>>> c.shape
TensorShape([10, 50, 20])

according to tensorflow's release log, RegisterShape is deprecated, and you should use SetShapeFn to define the shape when registering your operator in c++ source file.