How to apply control flow in nn.Module attributes?

[wztdream]

Hi,

It seems flax not allow modify Module attributes in setup. As an example, how to right below code in flax?

class Foo(nn.Module):
    def __init__(self, a=4,b=None)
        self.a = a
        self.b = b
        if b is None:
            self.b = a //2 if a !=0 else 1

I tried setting them in setup but trigger SetAttributeInModuleSetupError

class Foo(nn.Module):
    a:int=4
    b:Any=None
    def setup(self):
        if self.b is None:
            self.b = self.a//2 if self.a!=0 else 1

[wztdream]

Sorry, wrong place, It should be in flax repo. How to close or cancel it here?