RMSNorm Pallas Kernel

[AakashKumarNain]

I see that we have the new pallas rms norm layer as defined here. There is no documentation for the same in the docs. Can someone please elaborate the input shape it accepts i.e. Does it work for a single example, or does it work with batched inputs?

[superbobry]

The GPU kernels are indeed undocumented, but you could get some idea about the shape expectations from the tests.

[AakashKumarNain]

Perfect! Thanks