diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index c62b1cf3e23c..64148bb2414e 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -475,8 +475,8 @@ def log_softmax(data, axis=-1): data: tvm.relay.Expr The input data to the operator. - axis: int - The axis to sum over when computing softmax + axis: int, optional + The axis to sum over when computing log softmax Returns ------- @@ -510,6 +510,9 @@ def max_pool1d(data, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : int or tuple of int, optional The strides of pooling. @@ -567,6 +570,9 @@ def max_pool2d(data, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : tuple of int, optional The strides of pooling. @@ -612,6 +618,9 @@ def max_pool3d(data, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : tuple of int, optional The strides of pooling. @@ -658,6 +667,9 @@ def avg_pool1d(data, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : int or tuple of int, optional The strides of pooling. @@ -720,6 +732,9 @@ def avg_pool2d(data, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : tuple of int, optional The strides of pooling. @@ -769,6 +784,9 @@ def avg_pool3d(data, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : tuple of int, optional The strides of pooling. @@ -811,6 +829,9 @@ def max_pool2d_grad(out_grad, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : tuple of int, optional The strides of pooling. @@ -851,6 +872,9 @@ def avg_pool2d_grad(out_grad, data : tvm.relay.Expr The input data to the operator. + pool_size : int or tuple of int, optional + The size of window for pooling. + strides : tuple of int, optional The strides of pooling.