nn.compact vs setup for holoscope

[schneidermc]

In the holoscope example, trying to get the psf of the Holoscope class via
psf = holoscope.apply({"params": params, "state": state}, 0, method=holoscope.psf, rngs={"noise": key}).intensity.squeeze()
gives the following error:
AssignSubModuleError: Submodule ObjectivePointSource must be defined in 'setup()' or in a method wrapped in '@compact'

[diptodip]

Can you try with method=Holoscope.psf? But either way the Holoscope class in that example could definitely be improved if we make the psf method just use functions instead of elements. I'll get on that.