Training API with models that are not AbstractLuxLayers

[msainsburydale]

Thanks for your contributions to deep learning in Julia.

I have a package (NeuralEstimators.jl) that currently uses Flux. I'd like to make the package backend agnostic, in the sense that users can choose either Flux or Lux.

I'm nearly there, but I've run into problem that I'm finding difficult to solve. The public training API (e.g., TrainState, single_train_step!) requires an AbstractLuxLayer. However, declaring a struct to be an AbstractLuxLayer messes up the Flux training API (Flux.withgradient).

To simplify, I essentially have:

struct MyContainer
  field1
  field2
  ...
end

where the fields of MyContainer contain Flux models or Lux models. If the fields are Lux models, then I want to pass my_container::MyContainer into TrainState and single_train_step!.

Is this possible? For instance, could the type annotations on TrainState and single_train_step! be loosened, and we instead just define the necessary methods for our custom structs (e.g., initialparameters, initialstates)?

I think this is something that could come up with other packages in the ecosystem looking to adopt Lux while maintaining compatibility with Flux.

Thanks again for your work, and thanks in advance for any comments you might have.

[avik-pal]

Does https://lux.csail.mit.edu/stable/manual/flux_lux_interop work? We can try relaxing the types but I am pretty sure there are assumptions in codebase which will be violated. That said happy to review a PR if tests pass after relaxing the types

[msainsburydale]

Thanks for the suggestion! But I don't think that https://lux.csail.mit.edu/stable/manual/flux_lux_interop is what I'm looking for. On that page, it assumes that the models are Lux models, but I need to allow Flux models.

I got things working (see here if interested.) Basically, I just needed to:

1. remove the AbstractLuxLayer type annotation in TrainState
2. define a method preserves_state_type
3. define methods initialparameters/initialstates and parameterlength/statelength.

These are all currently given in my Lux extension (NeuralEstimator is the abstract type that I define methods for).

Essentially, I'm relying on duck typing, and it seems to work fine. I opened a pull request remove the AbstractLuxLayer type annotation in TrainState.

For Step 2, I wonder whether preserves_state_type could be determined automatically based on the fields of the container... the docs say that this is already done for an AbstractLuxContainerLayer, but I can't use AbstractLuxContainerLayer because it is a subtype of AbstractLuxLayer. Of course it isn't hard to just define preserves_state_type manually, but I think this would be a nice feature for user friendliness.

Generally speaking, I think it would be helpful if the package relied more on duck typing, particularly for codebases that want to maintain compatibility with Flux, though I'm not familiar with the codebase and this may not be straightforward.