Making sure parameters satisfy constraints

[ingmarschuster]

Hi,

How do I make sure the optimizer respects certain parameter constraints, for example positivity?

class GlobalScaler(ln.Module):
    scale_init: Callable[[PRNGKey, Shape, Dtype], Array] 

    @compact
    def __call__(self, inp):
        scale = self.param("scale", self.scale_init, (1,), np.float32, )
        return scale * inp

Now how do I make sure that the scale param is positive? I could use a bijection from R to R^+ of course. Is there a way to not hardcode which bijection is used?