jwfromm commented on code in PR #16654:
URL: https://github.com/apache/tvm/pull/16654#discussion_r1506388909
##########
python/tvm/relax/frontend/nn/op.py:
##########
@@ -413,10 +417,84 @@ def conv2d(
strides=stride,
padding=padding,
dilation=dilation,
+ data_layout=data_layout,
groups=groups,
)
if bias is not None:
- conv_out = _op.add(conv_out, _op.reshape(bias._expr, [1, -1, 1, 1]))
+ if data_layout == "NCHW":
+ conv_out = _op.add(conv_out, _op.reshape(bias._expr, [1, -1, 1,
1]))
+ elif data_layout == "NHWC":
+ conv_out = _op.add(conv_out, _op.reshape(bias._expr, [1, 1, 1,
-1]))
+ else:
+ raise NotImplementedError(f"Dont know how to handle layout
{data_layout}.")
+
+ return wrap_nested(conv_out, name)
+
+
+def conv3d(
+ x: Tensor,
+ weight: Tensor,
+ bias: Optional[Tensor] = None,
+ stride: Optional[Union[int, Tuple]] = 1,
+ padding: Optional[Union[int, Tuple, str]] = 0,
+ dilation: Optional[Union[int, Tuple]] = 1,
+ groups: Optional[int] = 1,
+ data_layout: Optional[str] = "NCDHW",
+ name: str = "conv3d",
+) -> Tensor:
+ """Applies a 3D convolution over an input image composed of sevaral input
planes
+
+ Parameters
+ ----------
+ x : Tensor
+ Input tensor of shape [B, N, D, H, W]
+
+ weight : Tensor
+ Filters of shape [O, N/groups, kD, kH, kW]
+
+ bias : Optional[Tensor]
+ Optional bias tensor of shape [O].
+
+ stride : Optional[Union[int, Tuple]]
+ The stride of the convolving kernel. Can be a single number
+ or tuple of (sD, sH, sW).
+
+ padding : Optional[[Union[int, Tuple]]]
+ Implicit paddings on both sides of the input.
+
+ dilation : Optional[Union[int, Tuple]]
+ The spacing between kernel elements. Can be a single number of tuple
(dD, dH, dW).
+
+ groups : Optional[int]
+ Split input into a number of groups.
+
+ data_layout : Optional[str]
+ Optional layout of the input and output data.
+
+ name : str
+ Name hint.
+
+ Returns
+ -------
+ result : Tensor
+ The computed result with shape [B, O, oD, oH, oW].
Review Comment:
I think the reason it uses O is to be consistent with weights, which almost
always use `OIHW` format. Other convolutions do the same. I agree something
like `oC` would be more legible but I'd like to stay consistent.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
