Currently there is no way to initialize the output tensor of a reduction with another Tensor.
One use case for such a functionality is when we have a fused convolution with bias-add operation. Currently define 2 compute operations, where the first one defines the reduction in the convolution, and another compute for bias add as shown below:
Output = te.compute(
(batch, out_height, out_width, out_channel),
lambda nn, yy, xx, ff: te.sum(
PaddedInput[nn, yy * stride_h + ry * dilation_h,
xx * stride_w + rx * dilation_w, rc].astype(out_dtype) *
Filter[ry, rx, rc, ff].astype(out_dtype), axis=[ry, rx, rc]),
name="Conv2dOutput", tag="conv2d_nhwc")
Output_with_bias = te.compute(
(batch, out_height, out_width, out_channel),
lambda nn, yy, xx, ff: Output[nn, yy, xx, ff] + Bias[nn, yy, xx, ff])
The code generated for the above computation would contain a fused loop nest, which has an initialization with 0f operation for the Output tensor, followed by the actual reduction for the convolution and then adding bias-add with Output. What I wanted to do is to remove the bias add computation and instead initialize the Output Tensor with the Bias values.
One way to do this would be to add an init member to ReduceNode object and an init argument to the reducer function returned by CommReducer, which can be used to pass in a Tensor with which to initialize the output Tensor instead of the identity element.
For Example, te.sum would be used as shown below:
te.sum(
PaddedInput[nn, yy * stride_h + ry * dilation_h,
xx * stride_w + rx * dilation_w, rc].astype(out_dtype) *
Filter[ry, rx, rc, ff].astype(out_dtype), axis=[ry, rx, rc],
init=Bias[nn, yy, xx, ff])
init=Bias[nn, yy, xx, ff] passes the Tensor that can be used in MakeReduction function in compute_op.cc to initialize with Bias instead of the identity element.
Let me know what you think about adding this option and any suggestions about my approach to do this.
Thanks,
Anirudh