Programmatically extend struct.dataclass class with additional attribute

[JohannesAck]

Hi,

I've been using flax for a few weeks now and I like it a lot!

I have a project where I would like to extend a flax.struct.dataclass received from external library code with an additional attribute for compatibility.

As a concrete example, I have a library function library_fn that gives me a dataclass with attributes a,b,c and would like to add to it an attribute d, e.g.

Of course this can be done manually for a single dataclass, but as I need to support a variety of dataclasses it would help a lot to do it programmatically.
I assumed this would be possible by declaring a new type and then manually applying the decorator struct.dataclass to it, but for some reason this raises an unexpected keyword argument:

import flax.struct as struct

@struct.dataclass
class LibraryDC:
    a: float = 0.0
    b: float = 2.0
    c: int = 1

def library_fn(x):
    return LibraryDC(a=x)

@struct.dataclass
class ManualWrapped(object):
    a: float = 0.0
    b: float = 2.0
    c: int = 1
    d: float = 0.0

def manually_wrapped_fn(x: LibraryDC):
    return ManualWrapped(d=0.0, **x.__dict__)

def wrapper_fn(x: LibraryDC):
    programmatic_wrapped_class = type('programmatic_wrapper_class', (object,), x.__dict__ | {'d': 0.0})
    wrapped_class = struct.dataclass(programmatic_wrapped_class)
    return wrapped_class(d=0.0, **x.__dict__)

if __name__ == '__main__':
    a = library_fn(1.0)
    print(a)  # LibraryDC(a=1.0, b=2.0, c=1)
    c = manually_wrapped_fn(a)
    print(c)  # ManualNew(a=1.0, b=2.0, c=1, d=0.0)
    b = wrapper_fn(a)
    print(b)

This outputs

LibraryDC(a=1.0, b=2.0, c=1)
ManualNew(a=1.0, b=2.0, c=1, d=0.0)
Traceback (most recent call last):
  File "/workspaces/test_file.py", line 36, in <module>
    b = wrapper_fn(a)
  File "/workspaces/test_file.py", line 27, in wrapper_fn
    return wrapped_class(d=0.0, **x.__dict__)
TypeError: inner_wrapper_class.__init__() got an unexpected keyword argument 'd'

I couldn't find a similar discussion anywhere, soI would be very grateful for any help with this issue!

[ASEM000]

this should work

import flax.struct as struct

@struct.dataclass
class LibraryDC:
    a: float = 0.0
    b: float = 2.0
    c: int = 1

def library_fn(x):
    return LibraryDC(a=x)

@struct.dataclass
class ManualWrapped(object):
    a: float = 0.0
    b: float = 2.0
    c: int = 1
    d: float = 0.0

def manually_wrapped_fn(x: LibraryDC):
    return ManualWrapped(d=0.0, **x.__dict__)

def wrapper_fn(x: LibraryDC):
    variables =x.__dict__ | {'d': 0.0}
+    annotations = {k: type(v) for k, v in variables.items()}
    programmatic_wrapped_class = type('programmatic_wrapper_class', (object,), {"__annotations__": annotations, **variables})
    wrapped_class = struct.dataclass(programmatic_wrapped_class)
    return wrapped_class(d=0.0, **x.__dict__)

if __name__ == '__main__':
    a = library_fn(1.0)
    print(a)  # LibraryDC(a=1.0, b=2.0, c=1)
    c = manually_wrapped_fn(a)
    print(c)  # ManualNew(a=1.0, b=2.0, c=1, d=0.0)
    b = wrapper_fn(a)
    print(b)
LibraryDC(a=1.0, b=2.0, c=1)
ManualWrapped(a=1.0, b=2.0, c=1, d=0.0)
programmatic_wrapper_class(a=1.0, b=2.0, c=1, d=0.0)

AFAIK , dataclasses looks at the class type hints defined by __annotations__ to construct the init method.
LMK if this helps