How to implement an Atan2 custom op?

[tkkhuu]

I am trying to implement a conversion for tf.Atan2 and hoping to follow the implementation of Atan, however it is blank in tensorflow-onnx/tf2onnx/onnx_opset/math.py:

@tf_op(["Acos", "Asin", "Atan", "Cos", "Sin", "Tan"])
class TrigOpSinceOpset7:
    @classmethod
    def version_7(cls, ctx, node, **kwargs):
        pass

Does this mean I can also do a similar thing to Atan2, something like this:

@tf_op("Atan2")
class MyAtan2:
    @classmethod
    def version_12(cls, ctx, node, **kwargs):
        pass

or do I have to implement it using the existing layers?

[guschmue]

onnx doesn't have atan2 but we can maybe compose it with a couple of ops.

[tkkhuu]

Thanks for your reply, I have attempted implementing using the existing ops:

@tf_op('MyAtan2')
class MyAtan2:

    @classmethod
    def version_1(cls, ctx, node, **kwargs):

        x = node.input[1]
        y = node.input[0]

        y_div_x_node = ctx.make_node('Div', [y, x], op_name_scope=node.name, name='y_div_x_node')
        atan_node = ctx.make_node('Atan', y_div_x_node.output, op_name_scope=node.name, name='atan_node')

        is_x_pos_node = ctx.make_node('Greater', [x, null], op_name_scope=node.name, name='is_x_pos_node')
        is_x_neg_node = ctx.make_node('Less', [x, null], op_name_scope=node.name, name='is_x_neg_node')
        is_x_zero_node = ctx.make_node('Equal', [x, null], op_name_scope=node.name, name='is_x_zero_node')

        is_y_pos_node = ctx.make_node('Greater', [y, null], op_name_scope=node.name, name='is_y_pos_node')
        is_y_neg_node = ctx.make_node('Less', [y, null], op_name_scope=node.name, name='is_y_neg_node')
        is_y_zero_node = ctx.make_node('Equal', [y, null], op_name_scope=node.name, name='is_y_zero_node')

But I'm not sure how to handle the ifs statement to condition x and y (e.g: positive or negative)

[guschmue]

I think you'd need to come up with a new tensor to multiply the results with.
But its tricky because there are multiple cases and you can't use 'If' (since it is not element wise).
Need to think a little about it.
There is some code that does a similar thing to implement Erf before there was a Erf op here:
https://github.com/onnx/tensorflow-onnx/blob/master/tf2onnx/onnx_opset/math.py#L421
but you case is worse because of the multiple conditions.

[tkkhuu]

I see, thank you for your help, yeah that's the example I'm using to create Atan2

[TomWildenhain-Microsoft]

The ifs can be handled by using the WHERE op if you are on opset >= 9

[TomWildenhain-Microsoft]

And the definitions here may be helpful: https://en.wikipedia.org/wiki/Atan2#Definition_and_computation

[guschmue]

yes, 'where' will do.