[REQ] Allow overriding __init__ of struct.dataclass

[PhilipVinc]

Python data classes allow the ovverriding of frozen data classes __init__ methods, that essentially just set all the relevant fields correctly.

I would like to ask that you expose that kwargument in your data class wrapper.

This is useful if a frozen data class is usually constructed starting from object_A but does not store object_A itself, but rather something computed starting from it.
We can get around this by using a constructor method, but for a series of reasons the easiest things for us to do in NetKet would be to just override the __init__.

I don't see reasons not to allow this, as it is allowed by the standard dataclass specification.

[PhilipVinc]

To give a more concrete example, we have a series of LinearOperator-like objects called FMJacobian and FMOnTheFly and FMJacobianPyTree (for example) which are lazy encodings of the fisher matrix (~hessian/natural gradient).

From an high level API they all take as input a model, the training data and the parameters, to be constructed.
However FMJacobian computes the jax.grad(model, pars, data) and ravels it to a dense array at construction and stores it. FMJacobianPyTree stores the gradient. FMOnTheFly stores model/data/parameters and precomputes nothing...

Ideally we'd do this in the __init__ method.
Again, we can use a factory method, but we would rather not because, among over reasons...

• We need to remember in type annotations to annotate with FMJacobianT but construct it with FMJacobian.
• we use annotations together with (multiple)dispatch to change behaviour. Having different names is bug prone and annoying.

[jheek]

This question pops up quite often but there is a major issue. A smart __init__ method (that is one where args don't map 1:1 to attributes) typically don't support clone & replace properly.

So in practice:
For a functional dataclass we want to update fields by calling cls(**dataclass_fields_as_dict(instance), **updates). This is only well defined if init simply maps the constructor args to attributes.

[PhilipVinc]

Ah, makes sense. Though it remains annoying.

I guess you need to flatten/unflatten the data class, right?
Because an alternative would be for you to add a custom flatten/unflatten_init function.

[PhilipVinc]

cc @femtobit @attila-i-szabo with whom we were discussing this

[jheek]

You can always define a vanilla dataclass with a custom flatten/unflatten. But in this case it's your own responsibility to reconstruct without a trivial constructor(*fields) call.

Disclaimer I'm a big Rust fanboy:
But I think "smart" constructors are not a good pattern. IMO a constructor should always be simple and what you want instead is a classmethod that creates your dataclass from some other value.
A simple argument for this is that once the constructor starts being smart you could imagine wanting multiple constructors. But to maintain sanity these constructors should probably have a self-explanatory name.
So I think what you should do is something like: FMJacobian.from_model(model, params, data)

The one exception to this could be simple type transformations that operate like a fixpoint. So for example transforming int -> (int,) such that when I pass in the (int,) again it stays (int,).

[PhilipVinc]

Disclaimer: I'm a Julia fanboy.
I like when interfaces are intuitive and simple.
In Julia, you can attach new definitions through multiple dispatch to the constructor of a class (it's a actually a function so there's nothing different from other cases).
A standard pattern is to build a dispatch ladder of constructors converting all inputs to the actual objects you need to actually build the class/structure.
so MyObject(x,y,z) always work, regardless of what you throw at it, and you don't need to define build_myobject_from_thatobject(that object).

Probably I like this pattern too much, as I try to replicate it in python. I am convinced that this leads to very natural-looking apis, because you never know what logic designers chose when they define those methods .from_model or .frommodel or .build_from_model... Clearly if you call MyClass(model) you want to construct your class from the model.

But that's a fight between patterns.

I actually realized that defining an __init__ automatically overrides dataclasses's __init__, as documented here, and this works also for flax.struct.dataclass.

However, your point is valid, and it should be used with care.

Sorry for the noise.

[jheek]

Disclaimer: I'm a Julia fanboy.

I haven't played with Julia enough. The multi-dispatch syntax is definitely more lightweight.
If Python had first-class multi-dispatch with ambiguity checking I agree your pattern works great!
Secretly I hope Julia just overtakes Python :)

[ToucheSir]

Just to build on the Julia point, here's the pytree-like library Flux uses for model (de)construction: https://github.com/FluxML/Functors.jl. It's very similar (e.g. fmap ≈ tree_multimap), but the biggest difference I've seen is that most libraries will internalize the tree mapping and rely on dispatch to process each leaf node instead of the Optax model of calling tree_multimap in every transform.