From 99b986885997e578c0f04cb9975778a800012b08 Mon Sep 17 00:00:00 2001
From: charan-003 <85248228+charan-003@users.noreply.github.com>
Date: Fri, 25 Apr 2025 13:18:30 -0600
Subject: [PATCH] Clarify dimension notation in conv1d, conv2d, and conv3d
 docstrings (#2123)

* Clarify dimension notation in conv1d, conv2d, and conv3d docstrings

* Updating transposed convs in conv1d, conv2d, and conv3d

---------

Co-authored-by: Sai Charan Arvapally <saicharan@Sais-MacBook-Pro.local>
---
 python/src/ops.cpp | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/python/src/ops.cpp b/python/src/ops.cpp
index 60b6188ed..a1e77d681 100644
--- a/python/src/ops.cpp
+++ b/python/src/ops.cpp
@@ -3455,8 +3455,8 @@ void init_ops(nb::module_& m) {
         1D convolution over an input with several channels
 
         Args:
-            input (array): Input array of shape ``(N, H, C_in)``.
-            weight (array): Weight array of shape ``(C_out, H, C_in)``.
+            input (array): Input array of shape ``(N, L, C_in)``.
+            weight (array): Weight array of shape ``(C_out, K, C_in)``.
             stride (int, optional): Kernel stride. Default: ``1``.
             padding (int, optional): Input padding. Default: ``0``.
             dilation (int, optional): Kernel dilation. Default: ``1``.
@@ -3514,7 +3514,7 @@ void init_ops(nb::module_& m) {
 
         Args:
             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
                 kernel strides. All spatial dimensions get the same stride if
                 only one number is specified. Default: ``1``.
@@ -3586,7 +3586,7 @@ void init_ops(nb::module_& m) {
 
         Args:
             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
                 kernel strides. All spatial dimensions get the same stride if
                 only one number is specified. Default: ``1``.
@@ -3619,8 +3619,8 @@ void init_ops(nb::module_& m) {
         1D transposed convolution over an input with several channels
 
         Args:
-            input (array): Input array of shape ``(N, H, C_in)``.
-            weight (array): Weight array of shape ``(C_out, H, C_in)``.
+            input (array): Input array of shape ``(N, L, C_in)``.
+            weight (array): Weight array of shape ``(C_out, K, C_in)``.
             stride (int, optional): Kernel stride. Default: ``1``.
             padding (int, optional): Input padding. Default: ``0``.
             dilation (int, optional): Kernel dilation. Default: ``1``.
@@ -3697,7 +3697,7 @@ void init_ops(nb::module_& m) {
 
         Args:
             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
                 kernel strides. All spatial dimensions get the same stride if
                 only one number is specified. Default: ``1``.
@@ -3783,7 +3783,7 @@ void init_ops(nb::module_& m) {
 
         Args:
             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
                 kernel strides. All spatial dimensions get the same stride if
                 only one number is specified. Default: ``1``.