google-deepmind / dm-haiku

JAX-based neural network library
https://dm-haiku.readthedocs.io
Apache License 2.0
2.91k stars 231 forks source link

Haiku initialization #606

Open adhikarirsr opened 1 year ago

adhikarirsr commented 1 year ago

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'))