Scalar parameters with constraints

[ingmarschuster]

I'm thinking of making jaxRK, a research library for advanced kernel methods, depend on linen for easing gradient based optimization.
What would be a recommended way to add scalar/array parameters with constraints? In a linen.Module subclass, I think using self.param is the way to declare the param, but what about the constraints?

[jheek]

If you want to define the parameter and the constraint together, you could define a special type for it like:

Array = Any
class Constrained(flax.struct.PyTreeNode):
  value: Array
  constraint: SomeType

That would allow you to later recover the parameters with their constraints.

[ingmarschuster]

cool. Would it also be possible to give an explicit parameter setting for initialization, in case a user doesn't want/need Gradient based optimization? E.g GaussianKernel(bandwidth=2.) is used for instantiating the GaussianKernel module with an explicit parameter, while the standard linen machinery is used in case of using optimization?

[jheek]

In the end variables are just a nested dictionary. But by convention we put parameters under the collection "params" indicating they are meant to be optimised. You can for example use self.variable('constants', 'some_name', init_fn) to make a separate collection of params that you don't want to optimize. The structure will look something like this

variables = Net.init(rng, ...)
print(variables)
{
  'params': {...}
  'constants': {...}
}

you can then easily optimize just the "params" collection and not the "constants" collection.

[ingmarschuster]

Thanks for your great answers. So I think I might write it up in the init method so that, for example, parameters that are provided to init might be saved as constants, while those that are not provided are params.