If I subclass nnx.Variable, then is there any functionality in NNX to validate the value provided to my subclass. For example, let's say I subclass into a Positive class for which the value provided must be positive

class PositiveReal(nnx.Variable):
    def __init__(self, value, **kwargs):
        # Apply a check here for value > 0
        super().__init__(value=jnp.asarray(value), **kwargs)

I could simply write assert jnp.all(value > 0); however, this is not jittable.