| # MIT License | |
| # | |
| # Copyright (c) 2026 audio-embeddings contributors | |
| # | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy | |
| # of this software and associated documentation files (the "Software"), to deal | |
| # in the Software without restriction, including without limitation the rights | |
| # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell | |
| # copies of the Software, and to permit persons to whom the Software is | |
| # furnished to do so, subject to the following conditions: | |
| # | |
| # The above copyright notice and this permission notice shall be included in all | |
| # copies or substantial portions of the Software. | |
| # | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | |
| # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | |
| # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | |
| # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | |
| # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | |
| # SOFTWARE. | |
| from __future__ import annotations | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class MixedPrecisionRMSNorm(nn.RMSNorm): | |
| """RMSNorm that keeps its master weight while matching activation dtype. | |
| PyTorch's fused RMSNorm requires the input and weight to have the same dtype. | |
| Mixed-precision training keeps parameters in FP32, so use a differentiable | |
| low-precision view of the weight for the operation itself. | |
| """ | |
| def forward(self, input: torch.Tensor) -> torch.Tensor: | |
| weight = self.weight | |
| if weight is not None and weight.dtype != input.dtype: | |
| weight = weight.to(dtype=input.dtype) | |
| return F.rms_norm(input, self.normalized_shape, weight, self.eps) | |