Forward autodiff : Multiplying by python float changes the dual dtype in some situations

ojoyot: <p>Hello</p>
<p>I have founded a strange behavior when I wanted to call func.jacfwd to autodiff one function. Here is a reduced example :</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print(torch.version.__version__)

fargs = {'dtype': torch.float32, 'device': 'cuda'}


A = torch.rand((5, 3, 100), **fargs)

def test(x):
    a = 0.5*x[0]
    p, t = torch.autograd.forward_ad.unpack_dual(a)
    print(f"a dtype : {p.dtype}, a.tangent dtype : {t.dtype}")
    v = torch.stack([a, *x[1:]])
    p, t = torch.autograd.forward_ad.unpack_dual(v)
    print(f"v dtype : {p.dtype}, v.tangent dtype : {t.dtype}")
    return torch.tensordot(v, A, dims=1)

torch.func.jacfwd(test)(torch.rand(5, **fargs))
</code></pre>
<p>This code gives the following output :</p>
<pre><code class="lang-auto">2.7.1+cu128
a dtype : torch.float32, a.tangent dtype : torch.float64
v dtype : torch.float32, v.tangent dtype : torch.float64    

Traceback (most recent call last):

[...]

return _VF.tensordot(a, b, dims_a, dims_b)  # type: ignore[attr-defined]
           ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: expected mat1 and mat2 to have the same dtype, but got: double != float
</code></pre>
<p>The line 10 (a = 0.5*x[0]) changes the dtype of the dual value but not the one of the primal (thus calling the function works but not calling its jacobian). This makes the code quite difficult to debug without knowing autograd.forward_ad.unpack_dual.</p>
<p>Two fixes I have founded :</p>
<ul>
<li>
<p>Replacing line 10 by <code>a = torch.tensor(0.5, **fargs)*x[0]</code> worked and is maybe the proper way to code.</p>
</li>
<li>
<p>Replacing line 10 by <code>a = (0.5*x)[0]</code>  worked surprisingly. The behavior is no longer the same and let me think that this might be a bug or an inconsistency ?</p>
</li>
</ul>
<pre><code class="lang-auto">2.7.1+cu128
a dtype : torch.float32, a.tangent dtype : torch.float32
v dtype : torch.float32, v.tangent dtype : torch.float32
</code></pre>

KFrank: <p>Hi Ojoyot!!</p>
<aside class="quote no-group quote-modified" data-username="ojoyot" data-post="1" data-topic="222163">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ojoyot/48/77902_2.png" class="avatar"> ojoyot:</div>
<blockquote>
<p>I have founded a strange behavior when I wanted to call func.jacfwd to autodiff one function.</p>
<p>…</p>
<p>This code gives the following output :</p>
<pre><code class="lang-auto">2.7.1+cu128
a dtype : torch.float32, a.tangent dtype : torch.float64
v dtype : torch.float32, v.tangent dtype : torch.float64    

</code></pre>
</blockquote>
</aside>
<p>I can reproduce your issue with the script you posted.</p>
<p>It looks like a bug to me.  (<a class="mention" href="/u/alband">@albanD</a> ?)  You might consider filling a github issue.</p>
<p>(As an aside, it happens with cpu tensors, as well as cuda.)</p>
<p>It appears to me that it is the combination of the python scalar multiplication</p>
<p>together with the forward-mode tensor being a 0-d tensor.  Here is a reproducing</p>
<p>script that illustrates this:</p>
<pre><code class="lang-auto">import torch
print (torch.__version__)

primal1d = torch.ones (1)       # 1-d tensor
tangent1d = torch.ones (1)      # 1-d tensor

primal0d = torch.ones (1)[0]    # 0-d tensor
tangent0d = torch.ones (1)[0]   # 0-d tensor

with torch.autograd.forward_ad.dual_level():
    print ('1-d tensor, no scalar multiplication -- ok')
    dual1d = torch.autograd.forward_ad.make_dual (primal1d, tangent1d)
    print ('dual1d.dtype:', dual1d.dtype, ', dual1d.shape:', dual1d.shape)
    p1d, t1d = torch.autograd.forward_ad.unpack_dual (dual1d)
    print ('p1d.dtype:', p1d.dtype, ', t1d.dtype:', t1d.dtype)
    print ('t1d.dtype == p1d.dtype:', t1d.dtype == p1d.dtype)
    
    print ('1-d tensor multiplied by python scalar -- ok')
    dual1db = 1.0 * dual1d
    print ('dual1db.dtype:', dual1db.dtype, ', dual1db.shape:', dual1db.shape)
    p1db, t1db = torch.autograd.forward_ad.unpack_dual (dual1db)
    print ('p1db.dtype:', p1db.dtype, ', t1db.dtype:', t1db.dtype)
    print ('t1db.dtype == p1db.dtype:', t1db.dtype == p1db.dtype)

    print ('0-d tensor, no scalar multiplication -- ok')
    dual0d = torch.autograd.forward_ad.make_dual (primal0d, tangent0d)
    print ('dual0d.dtype:', dual0d.dtype, ', dual0d.shape:', dual0d.shape)
    p0d, t0d = torch.autograd.forward_ad.unpack_dual (dual0d)
    print ('p0d.dtype:', p0d.dtype, ', t0d.dtype:', t0d.dtype)
    print ('t0d.dtype == p0d.dtype:', t0d.dtype == p0d.dtype)
    
    print ('0-d tensor multiplied by python scalar -- bad tangent dtype')
    dual0db = 1.0 * dual0d
    print ('dual0db.dtype:', dual0db.dtype, ', dual0db.shape:', dual0db.shape)
    p0db, t0db = torch.autograd.forward_ad.unpack_dual (dual0db)
    print ('p0db.dtype:', p0db.dtype, ', t0db.dtype:', t0db.dtype)
    print ('t0db.dtype == p0db.dtype:', t0db.dtype == p0db.dtype)

</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-auto">2.7.1+cu128
1-d tensor, no scalar multiplication -- ok
dual1d.dtype: torch.float32 , dual1d.shape: torch.Size([1])
p1d.dtype: torch.float32 , t1d.dtype: torch.float32
t1d.dtype == p1d.dtype: True
1-d tensor multiplied by python scalar -- ok
dual1db.dtype: torch.float32 , dual1db.shape: torch.Size([1])
p1db.dtype: torch.float32 , t1db.dtype: torch.float32
t1db.dtype == p1db.dtype: True
0-d tensor, no scalar multiplication -- ok
dual0d.dtype: torch.float32 , dual0d.shape: torch.Size([])
p0d.dtype: torch.float32 , t0d.dtype: torch.float32
t0d.dtype == p0d.dtype: True
0-d tensor multiplied by python scalar -- bad tangent dtype
dual0db.dtype: torch.float32 , dual0db.shape: torch.Size([])
p0db.dtype: torch.float32 , t0db.dtype: torch.float64
t0db.dtype == p0db.dtype: False

</code></pre>
<p>Best.</p>
<p>K. Frank</p>

albanD: <p>Thanks for the investigation <a class="mention" href="/u/kfrank">@KFrank</a>  !<br>
It definitely looks like a bug, the dtypes should match at all tie.</p>
<p>I guess the scalar dual is of the wrong dtype.</p>
<p>Would you be able to open an issue to track fixing this please?</p>

KFrank: <p>Hi Alban!</p>
<aside class="quote no-group" data-username="albanD" data-post="3" data-topic="222163">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>Would you be able to open an issue to track fixing this please?</p>
</blockquote>
</aside>
<p>My apologies, but I’m not able to create github issues at the moment.</p>
<p>(Google corrupted my account access and I haven’t set up a new account yet.)</p>
<p>Best.</p>
<p>K. Frank</p>

ojoyot: <p>Hi !</p>
<p>I have oppened a github issue at <a href="https://github.com/pytorch/pytorch/issues/160513" class="inline-onebox" rel="noopener nofollow ugc">Forward autodiff : Multiplying by python float changes the dual dtype in some situations · Issue #160513 · pytorch/pytorch · GitHub</a> .</p>
<p>Best</p>