I was wondering about the process for adding a new operator to TOPI. In a previous post, I remarked on the absence of a bias_add operator in TOPI when it seemed to already exist in Relay. This absence came up as a result of trying to prototype operators/schedules for a new hardware backend.
Since bias_add did not exist in TOPI, I implemented the op, and was wondering about how to go about upstreaming that op. Is there a suite of tests where this op should be added to verify functionality/stability? And are there any specific format guidelines for defining new operators?
I’ve attached the bias_add operator below:
@tvm.tag_scope(tag=tag.INJECTIVE)
def bias_add(x, b, layout=None):
"""Add biases to each feature map along channels dim.
Parameters
----------
x : tvm.Tensor
Input tensor. 4D
b : tvm.Tensor
Biases tensor. 1D = Channels of x
layout: string
layout of data
Returns
-------
y : tvm.Tensor
The result.
"""
# Infer layout from input layout
if layout is None:
layout = 'none'
if 'nchw' in x.op.tag: layout = 'NCHW'
if 'nhwc' in x.op.tag: layout = 'NHWC'
if layout == "NCHW":
batches, channels, height, width = x.shape
func = lambda bb, cc, hh, ww: tvm.sum(
x[bb, cc, hh, ww] + b[cc],
tvm.const(0, x.dtype))
elif layout == "NHWC":
batches, height, width, channels = x.shape
func = lambda bb, hh, ww, cc: tvm.sum(
x[bb, hh, ww, cc] + b[cc],
tvm.const(0, x.dtype))
else:
raise ValueError("bias_add does not support this layout {} yet".format(layout))
# Ensure we have enough biases
biases = b.shape
assert len(biases) == 1
biases = biases[0]
assert int(biases) == int(channels)
return tvm.compute(
x.shape, func)