native_layer_norm (for width dim) (#3001)
Summary: Pull Request resolved: https://github.com/pytorch/executorch/pull/3001 We implement `native_layer_norm` which has 3 outputs - normalization of the input tensor according to the given `normalized_shape` - mean - 1/sqrt(var + eps) ``` func: native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor) ``` According to SS-JIA's suggestion, a model specific implementation is more performant and preferred to a generic one. So we implemented the op in the following optimized way - our current use case has `normalized_shape` of len 1, namely we do the normalization through computing the mean and var at the last width dim - we do the computation in just one shader `native_layer_norm.glsl` without invoking the shaders to compute mean and var respectively - we use [Welford's online algorithm](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm) to compute mean and variance in one pass Reviewed By: SS-JIA, jorgep31415 Differential Revision: D56005629 fbshipit-source-id: 096c2e2f04b95f1f5c9205c4827091169771978c
W
Wei Lu committed
74576e83ddc4995baa4b544178f68df8a26db4a6
Parent: 075fe40
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
on 4/15/2024, 9:10:57 PM