May the methods in the torch.nn.functional of Pytorch be replaced in Burn in some way?

[AutumnRoom]

May the methods in the torch.nn.functional of Pytorch be replaced in Burn in some way?

Pytorch:
class LayerNorm(nn.Module):
def init(self, channels, eps=1e-5):
super().init()
self.channels = channels
self.eps = eps

self.gamma = nn.Parameter(torch.ones(channels))
self.beta = nn.Parameter(torch.zeros(channels))

def forward(self, x):
x = x.transpose(1, -1)
x = functional.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
return x.transpose(1, -1)

"x = functional.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)" can be replace?

[laggui]

We have some functional module ops here: https://github.com/tracel-ai/burn/blob/main/crates/burn-tensor/src/tensor/module.rs

But LayerNorm (and other normalization layers) don't have a functional form defined. Only as a module at this time.

[AutumnRoom]

It seems like that it can work in the same way.Trying to def a new "fn" may be no meaning but tough.

[laggui]

I guess we could expand the module functions to use with the layers, similar to convolutions atm. Feel free to open a PR if you think this should be done 🙂