class FourierFeatures(nn.Module):
    def __init__(self, start: int = 6, stop: int = 8, step: int = 1):
        super().__init__()

        self.start = start
        self.stop = stop
        self.step = step

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        r"""Forward method of the `FourierFeatures` class."""
        original_dtype = inputs.dtype
        inputs = inputs.to(torch.float32)
        num_channels = inputs.shape[1]
        num_freqs = (self.stop - self.start) // self.step

        freqs = torch.arange(self.start, self.stop, self.step, dtype=inputs.dtype, device=inputs.device)
        w = torch.pow(2.0, freqs) * (2 * torch.pi)  # [num_freqs]
        w = w.repeat(num_channels)[None, :, None, None, None]  # [1, num_channels * num_freqs, 1, 1, 1]

        # Interleaved repeat of input channels to match w
        h = inputs.repeat_interleave(num_freqs, dim=1)  # [B, C * num_freqs, T, H, W]
        # Scale channels by frequency.
        h = w * h

        return torch.cat([inputs, torch.sin(h), torch.cos(h)], dim=1).to(original_dtype)
/Users/nielsrogge/Documents/python_projecten/diffusers/src/diffusers/models/autoencoders/autoencoder_kl_mochi.py