Conflict between `dataclass` and `nn.Module`

username: Zephyrus
cooked: <p>When I decorate a module class <code>A</code> that inherits <code>nn.Module</code> with <code>dataclass</code>, the main model class <code>B</code> (class <code>A</code> is the sub module) that also inherits <code>nn.Module</code> will report an error when calling the <code>to(device)</code> method. The error message is <code>TypeError: unhashable type: 'A'</code>. This problem can be reproduced in the following simple way. are there some possible ways to solve this issue?</p>
<pre data-code-wrap="python"><code class="lang-python">@dataclass
class A(nn.Module):
    a: int
    b: int
    c: int

    def forward(self, x):
        return self.a * x + self.b * x + self.c
    
class B(nn.Module):
    def __init__(self, a, b, c):
        super().__init__()
        self.a = a
        self.b = b
        self.c = c
        self.sub_module = A(a, b, c)

    def forward(self, x):
        return self.a * x + self.b * x + self.c
    
model = B(1, 2, 3).to(torch.device('cuda'))
x = torch.tensor([1.0])
print(model(x))
</code></pre>

username: soulitzer
cooked: <p>You can get this to work with:</p>
<ul>
<li>pass unsafe_hash=True to the dataclass decorator or override it</li>
<li>overriding the default init that dataclass generates to also call super init</li>
</ul>
<p>It’s not super clear why you’d use dataclass though at this point since you’d have to override a lot of the default behavior anyway hash/init.</p>
<p>Also you might not even need to have A inherit nn.Module either since it appears to be stateless here.</p>

username: Zephyrus
cooked: <p>I also use JAX, in which <code>dataclass</code> is very convenient and can be registered as a <code>pytree</code> to leverage <code>autograd</code>. I just wonder if PyTorch has adapted to <code>dataclass</code>.</p>

username: soulitzer
cooked: <aside class="quote no-group" data-username="Zephyrus" data-post="3" data-topic="214355">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zephyrus/48/53193_2.png" class="avatar"> Zephyrus:</div>
<blockquote>
<p>I also use JAX, in which <code>dataclass</code> is very convenient and can be registered as a <code>pytree</code> to leverage <code>autograd</code>.</p>
</blockquote>
</aside>
<p>I see, thanks for the context.</p>
<p>Autograd doesn’t have extensive pytree support like jax, yes, and part of it is for bc-reasons. it is easy to write wrappers that do the pytree handling though.</p>
<blockquote>
<p>I just wonder if PyTorch has adapted to <code>dataclass</code> .</p>
</blockquote>
<p>For the purpose of autograd, I feel like I still don’t completely understand why it benefits to make a module a dataclass, bc you wouldn’t be passing that module instance itself to autograd right?</p>

username: Zephyrus
cooked: <aside class="quote no-group" data-username="soulitzer" data-post="4" data-topic="214355">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/48.png" class="avatar"> soulitzer:</div>
<blockquote>
<p>For the purpose of autograd, I feel like I still don’t completely understand why it benefits to make a module a dataclass, bc you wouldn’t be passing that module instance itself to autograd right?</p>
</blockquote>
</aside>
<p>Yes, as you said, I checked some information and it’s true that you shouldn’t pass a dataclass based module directly into the model, even in jax, it’s just to utilize the dataclass to better manage the model parameters. I got confused with these concepts, thank you very much for your answer, it helped me understand the difference.</p>