google-ai-edge / ai-edge-torch

Supporting PyTorch models with the Google AI Edge TFLite runtime.
Apache License 2.0
278 stars 36 forks source link

Fix aten.native_group_norm NHWC rewriter #100

Closed chunnienc closed 1 month ago

chunnienc commented 1 month ago

The decomposition for aten.native_group_norm runs precision issue with large weight and bias. Therefore the NHWC rewriter is disabled (inserts additional transposes) for native_group_norm with weight and bias input tensors.

BUG=https://b.corp.google.com/issues/353734124, https://b.corp.google.com/issues/354780253