Include any neural network as a component of a staged model

[mlprt]

It should be straightforward to allow any Equinox-based neural network to be component of an AbstractStagedModel, such as SimpleFeedback, and be called during one of the model stages.

We would need to:

1. if it doesn't take a state/hidden argument (e.g. not an RNN), wrap it to ignore the state argument that AbstractStagedModel.__call__ will try to pass it,
2. associate it with a generic NetworkState-like PyTree that has a single leaf, which stores the output of the module. Alternatively, keep a single Array in the PyTree of any model of which the network is a component (e.g. SimpleFeedback).

On the other hand, we could specify the network itself as an AbstractStagedModel where network layers correspond to distinct stages, and where the activities of different layers may be kept as part of the state. In that case, the user can add interventions to these states without needing to redesign the model.

SimpleStagedNetwork is the prototype for a neural network AbstractStagedModel.

[mlprt]

The concern over associating third-party modules with states should probably be simpler if we replace the staged model approach with a DAG approach (#28).