Refactor out dataclass transform

[copybara-service[bot]]

Refactor out dataclass transform that allows parent and name to be moved to the end of the argument list into a more general "kw_only_dataclasses" module. This module provides wrappers for dataclasses.dataclass and dataclasses.field that simulate support for keyword-only fields for Python versions before 3.10 (which is the version where dataclasses added keyword-only field support). For version 3.10+, we can just use the standard kw_only support.

[cgarciae]

Hey @edloper, I was chatting with @levskaya about this PR, if we want to fully support keyword-only arguments for 3.10 dataclasses it might be worth looking into the KW_ONLY typehint that tries to behave like python's * separator symbol for function signatures:

@dataclass
class Point:
  x: float
  _: KW_ONLY
  y: float
  z: float

p = Point(0, y=1.5, z=2.0)

In this example y and z are set to keyword-only because the go after the KW_ONLY type annotation. Maybe this is an overkill but it is part of the spec so it depends on what we want.

[edloper]

Updated the changelist to support KW_ONLY tag. But note that there are some limitations to this -- it doesn't perfectly emulate the 3.10 dataclass kw_only support. In particular:

1. It only supports kw_only args that have default values.
2. It just moves the kw_only fields to the end of the signature; it doesn't actually make them keyword_only if you look at inspect.signature(MyDataclass.init).

On Thu, Sep 22, 2022 at 5:28 PM Cristian Garcia @.***> wrote:

Hey @edloper https://github.com/edloper, I was chatting with @levskaya https://github.com/levskaya and if we want to fully support keyword-only arguments for 3.10 dataclasses it might be worth looking into the KW_ONLY https://docs.python.org/3/library/dataclasses.html#dataclasses.KW_ONLY typehint that tries to behave like python's * separator symbol for function signatures:

@dataclassclass Point: x: float _: KW_ONLY y: float z: float p = Point(0, y=1.5, z=2.0)

In this example y and z are set to keyword-only because the go after the KW_ONLY type annotation. Maybe this is an overkill but it is part of the spec so it depends on what we want.

— Reply to this email directly, view it on GitHub https://github.com/google/flax/pull/2468#issuecomment-1255565822, or unsubscribe https://github.com/notifications/unsubscribe-auth/ABMFVDENZPYR3TGL3RIEZSLV7TFOFANCNFSM6AAAAAAQMZRRB4 . You are receiving this because you were mentioned.Message ID: @.***>

[codecov-commenter]

Codecov Report

Merging #2468 (0190918) into main (69163b9) will increase coverage by 0.12%. The diff coverage is 93.33%.

@@            Coverage Diff             @@
##             main    #2468      +/-   ##
==========================================
+ Coverage   78.88%   79.01%   +0.12%     
==========================================
  Files          48       49       +1     
  Lines        5120     5175      +55     
==========================================
+ Hits         4039     4089      +50     
- Misses       1081     1086       +5     

Impacted Files	Coverage Δ
flax/linen/kw_only_dataclasses.py	92.75% <92.75%> (ø)
flax/linen/module.py	92.63% <100.00%> (-0.18%)	:arrow_down:

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.