Allow passing networks directly instead of using network=str(), network_kwargs=dict()

[LarsKue]

This has been fundamentally limited by serialization, which should become much simpler with #342.

The idea is that we want to move from this:

inference_network = FlowMatching(subnet="mlp", subnet_kwargs=dict(widths=[256, 512, 256]))

to this:

inference_network = FlowMatching(subnet=MLP(widths=[256, 512, 256]))

Advantages:

• Intuitive code structure
• Significantly improved transparency on which arguments are available
• Trivially allows users to pass their own networks
• Reduced developer load for maintaining dispatch

Disadvantages:

• Slightly more complicated type hinting

[LarsKue]

Now relates to #391

[LarsKue]

We will further deprecate the _kwargs pattern.

[vpratz]

We will further deprecate the _kwargs pattern.

Could you please detail what you mean by this and what alternative pattern you have in mind?

[LarsKue]

@vpratz You can check the behavior on the allow-networks branch. It just means that when the user passes subnet_kwargs to a model where it is not required (e.g. FlowMatching, but not CouplingFlow), they will get a warning that they should instantiate the network directly themselves. Other _kwargs are currently unaffected, such as optimal_transport_kwargs.

[vpratz]

I see, thanks for clarifying :)