Add GroupNorm support

[FL33TW00D]

#131 requires GroupNorm support

[AmineDiro]

Hi @FL33TW00D I'd like to take a swing at implementing this layer. GroupNorm can be seen as a generalization of LayerNorm but it differs in implementation enough that I think we can add it to the enum and reimplement the kernel

pub enum Norm {
    LayerNorm(LayerNorm),
   GroupNorm(GroupNorm),
}

Quick question, is there a reason for specializing the LayerNorm for vec2 and vec4 separatly?

[FL33TW00D]

@AmineDiro

Enum sounds good!

In WebGPU you can often get big performance boosts from vectorizing, vec2 isn't much of a difference but vec4 can be considerable. Even just a scalar kernel is great!

[AmineDiro]

Hello, Finally found time to look at the implementation of GroupNorm. If I am not mistaken I think that we have everything in thenorm.rs but I need to check if I am not missing something. For a given tensor of shape B,C, IMG_SIZE applying a groupnorm of num_groups is equivalent to applying LayerNorm on the tensor as B,n_groups, n_channels*IMG_SIZE . Is this coherent with the memory layout of ratchet ?

[FL33TW00D]

Hey @AmineDiro,
Ratchet memory layout is the following:

1. All buffers are contiguous
2. All are laid out 1D row major

Everything else is syntax sugar!

[AmineDiro]

Great so I think it should work! I'll try to push something as soon as possible !