Skip to content

Commit

Permalink
Remove the uninitialized field "embedding" in nn.Embed
Browse files Browse the repository at this point in the history
The presence of this field breaks dataclasses.asdict() but it doesn't seem to serve much purpose. It was added in a change that tries to reduce gpylint errors but removing this line does not trigger any new gpylint error.

PiperOrigin-RevId: 620422049
  • Loading branch information
kho authored and Flax Authors committed Apr 1, 2024
1 parent cc740d4 commit 9a2a4d6
Showing 1 changed file with 0 additions and 3 deletions.
3 changes: 0 additions & 3 deletions flax/linen/linear.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@

"""Linear modules."""

import dataclasses
from typing import (
Any,
Iterable,
Expand Down Expand Up @@ -1105,8 +1104,6 @@ class Embed(Module):
param_dtype: Dtype = jnp.float32
embedding_init: Initializer = default_embed_init

embedding: Array = dataclasses.field(init=False)

def setup(self):
self.embedding = self.param(
'embedding',
Expand Down

0 comments on commit 9a2a4d6

Please sign in to comment.