Clarify dimension notation in conv1d, conv2d, and conv3d docstrings

This commit is contained in:
Sai Charan Arvapally 2025-04-25 02:01:52 -06:00
parent f0e70afff0
commit 268314cc2f

View File

@ -3455,8 +3455,8 @@ void init_ops(nb::module_& m) {
1D convolution over an input with several channels 1D convolution over an input with several channels
Args: Args:
input (array): Input array of shape ``(N, H, C_in)``. input (array): Input array of shape ``(N, L, C_in)``.
weight (array): Weight array of shape ``(C_out, H, C_in)``. weight (array): Weight array of shape ``(C_out, K, C_in)``.
stride (int, optional): Kernel stride. Default: ``1``. stride (int, optional): Kernel stride. Default: ``1``.
padding (int, optional): Input padding. Default: ``0``. padding (int, optional): Input padding. Default: ``0``.
dilation (int, optional): Kernel dilation. Default: ``1``. dilation (int, optional): Kernel dilation. Default: ``1``.
@ -3514,7 +3514,7 @@ void init_ops(nb::module_& m) {
Args: Args:
input (array): Input array of shape ``(N, H, W, C_in)``. input (array): Input array of shape ``(N, H, W, C_in)``.
weight (array): Weight array of shape ``(C_out, H, W, C_in)``. weight (array): Weight array of shape ``(C_out, KH, KW, C_in)``.
stride (int or tuple(int), optional): :obj:`tuple` of size 2 with stride (int or tuple(int), optional): :obj:`tuple` of size 2 with
kernel strides. All spatial dimensions get the same stride if kernel strides. All spatial dimensions get the same stride if
only one number is specified. Default: ``1``. only one number is specified. Default: ``1``.
@ -3586,7 +3586,7 @@ void init_ops(nb::module_& m) {
Args: Args:
input (array): Input array of shape ``(N, D, H, W, C_in)``. input (array): Input array of shape ``(N, D, H, W, C_in)``.
weight (array): Weight array of shape ``(C_out, D, H, W, C_in)``. weight (array): Weight array of shape ``(C_out, KD, KH, KW, C_in)``.
stride (int or tuple(int), optional): :obj:`tuple` of size 3 with stride (int or tuple(int), optional): :obj:`tuple` of size 3 with
kernel strides. All spatial dimensions get the same stride if kernel strides. All spatial dimensions get the same stride if
only one number is specified. Default: ``1``. only one number is specified. Default: ``1``.