<p dir="auto">Recently, <code class="notranslate">functional.softmax</code> has been changed to accept <code class="notranslate">dim</code> as argument.<br>
When run on CPU, I can specify arbitrary value for <code class="notranslate">dim</code>, whether <code class="notranslate">input</code> is 2D or 3D.</p>
<p dir="auto">For example, the following code works without any problem:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="In [23]: functional.softmax(Variable(torch.ones(3, 4)), dim=0)
Out[23]: 
Variable containing:
 0.3333  0.3333  0.3333  0.3333
 0.3333  0.3333  0.3333  0.3333
 0.3333  0.3333  0.3333  0.3333
[torch.FloatTensor of size 3x4]

In [24]: functional.softmax(Variable(torch.ones(3, 4, 5)), dim=1)
Out[24]: 
Variable containing:
(0 ,.,.) = 
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500

(1 ,.,.) = 
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500

(2 ,.,.) = 
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
[torch.FloatTensor of size 3x4x5]"><pre class="notranslate"><code class="notranslate">In [23]: functional.softmax(Variable(torch.ones(3, 4)), dim=0)
Out[23]: 
Variable containing:
 0.3333  0.3333  0.3333  0.3333
 0.3333  0.3333  0.3333  0.3333
 0.3333  0.3333  0.3333  0.3333
[torch.FloatTensor of size 3x4]

In [24]: functional.softmax(Variable(torch.ones(3, 4, 5)), dim=1)
Out[24]: 
Variable containing:
(0 ,.,.) = 
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500

(1 ,.,.) = 
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500

(2 ,.,.) = 
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
  0.2500  0.2500  0.2500  0.2500  0.2500
[torch.FloatTensor of size 3x4x5]
</code></pre></div>
<p dir="auto">However, when the computation is performed in GPU, the flexibility decreases.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="In [25]: functional.softmax(Variable(torch.ones(3, 4)).cuda(), dim=0)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-25-d95ff8a04b77&gt; in &lt;module&gt;()
----&gt; 1 functional.softmax(Variable(torch.ones(3, 4)).cuda(), dim=0)

~/anaconda3/envs/ml/lib/python3.6/site-packages/torch/nn/functional.py in softmax(input, dim, _stacklevel)
    648     if dim is None:
    649         dim = _get_softmax_dim('softmax', input.dim(), _stacklevel)
--&gt; 650     return _Softmax(dim)(input)
    651 
    652 

RuntimeError: invalid argument 4: dim has to be 1 for 2D input at /home/jhchoi/workspace/pytorch/torch/lib/THCUNN/generic/SoftMax.cu:32

In [26]: functional.softmax(Variable(torch.ones(3, 4, 5)).cuda(), dim=1)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-26-672d6d5d3b8d&gt; in &lt;module&gt;()
----&gt; 1 functional.softmax(Variable(torch.ones(3, 4, 5)).cuda(), dim=1)

~/anaconda3/envs/ml/lib/python3.6/site-packages/torch/nn/functional.py in softmax(input, dim, _stacklevel)
    648     if dim is None:
    649         dim = _get_softmax_dim('softmax', input.dim(), _stacklevel)
--&gt; 650     return _Softmax(dim)(input)
    651 
    652 

RuntimeError: invalid argument 4: dim has to be 0 for 3D input at /home/jhchoi/workspace/pytorch/torch/lib/THCUNN/generic/SoftMax.cu:42"><pre class="notranslate"><code class="notranslate">In [25]: functional.softmax(Variable(torch.ones(3, 4)).cuda(), dim=0)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-25-d95ff8a04b77&gt; in &lt;module&gt;()
----&gt; 1 functional.softmax(Variable(torch.ones(3, 4)).cuda(), dim=0)

~/anaconda3/envs/ml/lib/python3.6/site-packages/torch/nn/functional.py in softmax(input, dim, _stacklevel)
    648     if dim is None:
    649         dim = _get_softmax_dim('softmax', input.dim(), _stacklevel)
--&gt; 650     return _Softmax(dim)(input)
    651 
    652 

RuntimeError: invalid argument 4: dim has to be 1 for 2D input at /home/jhchoi/workspace/pytorch/torch/lib/THCUNN/generic/SoftMax.cu:32

In [26]: functional.softmax(Variable(torch.ones(3, 4, 5)).cuda(), dim=1)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-26-672d6d5d3b8d&gt; in &lt;module&gt;()
----&gt; 1 functional.softmax(Variable(torch.ones(3, 4, 5)).cuda(), dim=1)

~/anaconda3/envs/ml/lib/python3.6/site-packages/torch/nn/functional.py in softmax(input, dim, _stacklevel)
    648     if dim is None:
    649         dim = _get_softmax_dim('softmax', input.dim(), _stacklevel)
--&gt; 650     return _Softmax(dim)(input)
    651 
    652 

RuntimeError: invalid argument 4: dim has to be 0 for 3D input at /home/jhchoi/workspace/pytorch/torch/lib/THCUNN/generic/SoftMax.cu:42
</code></pre></div>
<p dir="auto">Is it intended behavior? If it is, then I think it should be specified in the document.</p>
<p dir="auto">(removed previous report addressed by <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="267805013" data-permission-text="Title is private" data-url="https://github.com/pytorch/pytorch/issues/3245" data-hovercard-type="pull_request" data-hovercard-url="/pytorch/pytorch/pull/3245/hovercard" href="https://github.com/pytorch/pytorch/pull/3245">#3245</a>)</p>
<p dir="auto">There's a missing letter 'f' in the <a href="http://pytorch.org/docs/master/nn.html?highlight=softmax#torch.nn.functional.softmax" rel="nofollow">docs</a>, my browser doesn't render the equation:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=":math:`softmax(x) =

rac{exp(-x_i)}{sum_j exp(-x_j)}`"><pre class="notranslate"><code class="notranslate">:math:`softmax(x) =

rac{exp(-x_i)}{sum_j exp(-x_j)}`
</code></pre></div>