Not sure if a bug or not...

[jamestwebber]

I opened #35 because I thought I spotted a bug in the weight standardization code but @vballoli says it's fine, so I'm opening a discussion (which I never knew existed!) to figure it out.

The code reads

fan_in = torch.prod(torch.tensor(self.weight.shape[0:]))

This was suspicious to me as shape[0:] is just making a needless copy of shape, and calculates fan_in as the size of the entire tensor.

The code in the deepmind repository here reads shape[:-1], which means fan_in is the product over all but the last dimension, which makes more sense to me.

Maybe I am missing a pytorch vs jax implementation difference? What's the reason for the discrepancy?

[vballoli]

Thanks for raising this! I misinterpreted the issue yesterday, sorry about that. I've fixed the fan-in. Do let me know if these implementations match.

[jamestwebber]

This is still not matching, it's now shape[1:] (skipping the first dim) but the original implementation is shape[:-1] (skipping the last dim).

If there's some implicit transpose elsewhere that I didn't see, maybe this makes sense, but I think it should be shape[:-1]

[vballoli]

I don't think the weight shape is same as the Jax. See https://github.com/rwightman/pytorch-image-models/blob/4ea593196414684d2074cbb81d762f3847738484/timm/models/layers/std_conv.py#L79.

[vballoli]

The iokk orientation is different afaik.

[jamestwebber]

ah! I always forget this. the torch weights are multiplied from the left which is never what I expect to happen.