refactor: simplifications in ModuleMeta __call__

[nstarman]

No description provided.

[nstarman]

@patrick-kidger I rebased on dev but the target is still main, which I don't believe I can switch without closing this PR and opening a new one. Hence the pickup of other commits.
Cool. Didn't know you could switch targets post PR creation.

[nstarman]

@patrick-kidger I think this is ready for review.

[patrick-kidger]

So for this one I'm actually planning on doing my Module refactor fairly shortly. (As my private testing with it has so far gone smoothly.) For that reason I'm inclined not to change the internals here, and instead tag you on the PR when I write it? It'd be great to have your feedback on whether the new version looks up to scratch.

[nstarman]

Some possible ideas, inspired from dataclassish

1. Always define a __dataclass_init__ method that is constructed by Module and which is the init method created by dataclass if init=True. This is very useful when defining a custom init method to be able to do

class MyClass
    def __init__(self, ...):
        ... # Whatever manual processing
        self.__dataclass_init__(...)  # don't have to worry about setting in frozen vs not, descriptors vs not, etc.

Of course it's totally optional to use, it's just very convenient!
One of my common gripes when writing an __init__ with a dataclass is that I just needed to do a bit of manual processing but now I have to write all the attribute assignments.

2. Always define a __converter_init__ method that is constructed by Module and which does

def __converter_init__(self, ...):
    ba = run_converters(...)
    self.__dataclass_init__(**ba)

When dataclass_init=True this is swapped in for __init__ so that input conversion happens early, which should be more efficient (e.g. saving 2x assignment).
When dataclass_init=False then users can simplify their __init__ methods by doing

class MyClass
    def __init__(self, ...):
        ... # Whatever manual processing
        self. __converter_init__(...)  # handles conversion, and all complexities of setting.

Of course it's totally optional to use, it's just very convenient!

3. If dataclass_init is True and the run_converters skip args that are Tracers then the __init__ method satisfies JAX's assumptions about dataclasses and we can use the faster pytree registration.

[patrick-kidger]

These are an interesting set of ideas! I can totally see why these are attractive for power users.

Points (1) and (2) I think probably impose too great a learning curve for new users: we already have converter, __post_init__, __init__ and __check_init__, which is already rather a lot. These methods might save a few lines but I think they increase the complexity too much.

Point (3) sounds like an internal optimisation we could probably perform, however. I like the sound of that one.

[nstarman]

SGTM. I think we can also do the converter loop as part of the __init__ trick when dataclass_init=True as an internal optimization. That should reduce setting from 2x to 1x.