google-deepmind/dm-haiku

Haiku initialization

adhikarirsr opened this issue · 0 comments

Pytorch initialization:

def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

Is the above initialization equivalent to :
hk.Linear(hidden_size, with_bias=True, w_init=VarianceScaling(1.0, 'fan_in', 'uniform'),b_init= VarianceScaling(1.0, 'fan_in', 'uniform'))