Feature Request: some way to pass hyperparameters out of transforms

[google-deepmind]

Feature Request: some way to pass hyperparameters out of transforms

[cgarciae]

Hey, this seems like a good use case for creating a custom hook, in Elegy we created a custom elegy.transform to enable custom hooks like elegy.add_loss and elegy.add_metric. I imagine you could create an API like this:

def forward(x):
    a = get_hyperparameter("a")
    b = get_hyperparameter("b")
    return MyModule(a, b)(x)

model = custom_transform(forward)

hyperparameters = dict(a=1, b=2)

model.init(..., hyperparameters)
model.apply(..., hyperparameters, ...)

But creating you own transform is not that straight forward, you can check our implementation here.

Now that you mention it I think this would be a good feature for Elegy. I it would also be nice if we could make a transform general enough so users could create their own hooks.