Bias initialization based on the input dimension of the layer

Hey,

I was trying to initialize the biases of my layers according the way it's done in [PyTorch's default way of doing so, for linear layers](https://github.com/pytorch/pytorch/blob/2400cb1d57aa596b503bd1c5c567113ed4962016/torch/nn/modules/linear.py#L108-L111):
```python
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
init.uniform_(self.bias, -bound, bound)
```
but I realized it's almost impossible to do so in a clean way using flax. As in an initializer, I only have access to the shape of the parameter being initialized, which for bias translates to the output dimension of a layer. I imagine this was intended (https://github.com/google/jax/issues/2075#issuecomment-578465814, https://github.com/google/flax/issues/2749) but I also imagine this could be annoying for a lot of people like me, who would like to initialize their biases based on the input dimension in some way (let's say to numerically test some theoretical analysis).

Is there any recommended way of initializing the biases like this? If not, is this supposed to be supported in the future at all?

Thanks!