reflectorch 1.3.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

reflectorch/models/encoders/fno.py

@@ -1,134 +1,134 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from reflectorch.models.activations import activation_by_name
-
-class SpectralConv1d(nn.Module):
-    def __init__(self, in_channels, out_channels, modes):
-        super().__init__()
-
-        """
-        1D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
-        """
-
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.modes = modes  #Number of Fourier modes to multiply, at most floor(N/2) + 1
-
-        self.scale = (1 / (in_channels*out_channels))
-        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, modes, dtype=torch.cfloat))
-
-    # Complex multiplication
-    def compl_mul1d(self, input, weights):
-        # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
-        return torch.einsum("bix,iox->box", input, weights)
-
-    def forward(self, x):
-        batchsize = x.shape[0]
-        #Compute Fourier coeffcients up to factor of e^(- something constant)
-        x_ft = torch.fft.rfft(x)
-
-        # Multiply relevant Fourier modes
-        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1,  device=x.device, dtype=torch.cfloat)
-        out_ft[:, :, :self.modes] = self.compl_mul1d(x_ft[:, :, :self.modes], self.weights1)
-
-        #Return to physical space
-        x = torch.fft.irfft(out_ft, n=x.size(-1))
-        return x
-    
-
-class FnoEncoder(nn.Module):
-    """An embedding network based on the Fourier Neural Operator (FNO) architecture
-
-    .. image:: ../documentation/fig_reflectometry_embedding_networks.png
-        :width: 400px
-        :align: center
-
-    Args:
-        in_channels (int): number of input channels
-        dim_embedding (int): dimension of the output embedding
-        modes (int): number of Fourier modes
-        width_fno (int): number of channels of the intermediate representations
-        n_fno_blocks (int): number of FNO blocks
-        activation (str): the activation function
-        fusion_self_attention (bool): whether to use fusion self attention for merging the tokens (instead of mean)
-        fsa_activation (str): the activation function of the fusion self attention block
-    """
-    def __init__(
-            self, 
-            in_channels: int = 2, 
-            dim_embedding: int = 128, 
-            modes: int = 32, 
-            width_fno: int = 64, 
-            n_fno_blocks: int = 6, 
-            activation: str = 'gelu',
-            fusion_self_attention: bool = False,
-            fsa_activation: str = 'tanh',
-            ):
-        super().__init__()
-
-
-        self.in_channels = in_channels
-        self.dim_embedding = dim_embedding
-        
-        self.modes = modes
-        self.width_fno = width_fno
-        self.n_fno_blocks = n_fno_blocks
-        self.activation = activation_by_name(activation)()
-        self.fusion_self_attention = fusion_self_attention
-        
-
-        self.fc0 = nn.Linear(in_channels, width_fno) #(r(q), q)
-        self.spectral_convs = nn.ModuleList([
-            SpectralConv1d(in_channels=width_fno, out_channels=width_fno, modes=modes) for _ in range(n_fno_blocks)
-            ])
-        self.w_convs = nn.ModuleList([
-            nn.Conv1d(in_channels=width_fno, out_channels=width_fno, kernel_size=1) for _ in range(n_fno_blocks)
-            ])
-        self.fc_out = nn.Linear(width_fno, dim_embedding)
-
-        if fusion_self_attention:
-            self.fusion = FusionSelfAttention(embed_dim=width_fno, hidden_dim=2*width_fno, activation=fsa_activation)
-        
-    def forward(self, x):
-        """"""
-
-        x = x.permute(0, 2, 1) #(B, D, S) -> (B, S, D)
-        x = self.fc0(x)
-        x = x.permute(0, 2, 1) #(B, S, D) -> (B, D, S) 
-
-        for i in range(self.n_fno_blocks):
-            x1 = self.spectral_convs[i](x)
-            x2 = self.w_convs[i](x)
-            
-            x = x1 + x2
-            x = self.activation(x)
-
-        if self.fusion_self_attention:
-            x = x.permute(0, 2, 1)
-            x = self.fusion(x)  
-        else:
-            x = x.mean(dim=-1)
-
-        x = self.fc_out(x)
-        
-        return x
-    
-
-class FusionSelfAttention(nn.Module):
-    def __init__(self, embed_dim: int = 64, hidden_dim: int = 64, activation: str = 'gelu'):
-        super().__init__()
-        activation = activation_by_name(activation)()
-        self.fuser = nn.Sequential(nn.Linear(embed_dim, hidden_dim), 
-                                   activation,
-                                   nn.Linear(hidden_dim, 1, bias=False))
-        
-    def forward(self, 
-                c: torch.Tensor,  # (batch_size x seq_len x embed_dim)
-                mask: torch.Tensor = None, # (batch_size x seq_len)
-                ):  
-        a = self.fuser(c)
-        alpha = torch.exp(a)*mask.unsqueeze(-1) if mask is not None else torch.exp(a)
-        alpha = alpha/alpha.sum(dim=1, keepdim=True)
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from reflectorch.models.activations import activation_by_name
+
+class SpectralConv1d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes):
+        super().__init__()
+
+        """
+        1D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
+        """
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes = modes  #Number of Fourier modes to multiply, at most floor(N/2) + 1
+
+        self.scale = (1 / (in_channels*out_channels))
+        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, modes, dtype=torch.cfloat))
+
+    # Complex multiplication
+    def compl_mul1d(self, input, weights):
+        # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+        return torch.einsum("bix,iox->box", input, weights)
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        #Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfft(x)
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-1)//2 + 1,  device=x.device, dtype=torch.cfloat)
+        out_ft[:, :, :self.modes] = self.compl_mul1d(x_ft[:, :, :self.modes], self.weights1)
+
+        #Return to physical space
+        x = torch.fft.irfft(out_ft, n=x.size(-1))
+        return x
+    
+
+class FnoEncoder(nn.Module):
+    """An embedding network based on the Fourier Neural Operator (FNO) architecture
+
+    .. image:: ../documentation/fig_reflectometry_embedding_networks.png
+        :width: 400px
+        :align: center
+
+    Args:
+        in_channels (int): number of input channels
+        dim_embedding (int): dimension of the output embedding
+        modes (int): number of Fourier modes
+        width_fno (int): number of channels of the intermediate representations
+        n_fno_blocks (int): number of FNO blocks
+        activation (str): the activation function
+        fusion_self_attention (bool): whether to use fusion self attention for merging the tokens (instead of mean)
+        fsa_activation (str): the activation function of the fusion self attention block
+    """
+    def __init__(
+            self, 
+            in_channels: int = 2, 
+            dim_embedding: int = 128, 
+            modes: int = 32, 
+            width_fno: int = 64, 
+            n_fno_blocks: int = 6, 
+            activation: str = 'gelu',
+            fusion_self_attention: bool = False,
+            fsa_activation: str = 'tanh',
+            ):
+        super().__init__()
+
+
+        self.in_channels = in_channels
+        self.dim_embedding = dim_embedding
+        
+        self.modes = modes
+        self.width_fno = width_fno
+        self.n_fno_blocks = n_fno_blocks
+        self.activation = activation_by_name(activation)()
+        self.fusion_self_attention = fusion_self_attention
+        
+
+        self.fc0 = nn.Linear(in_channels, width_fno) #(r(q), q)
+        self.spectral_convs = nn.ModuleList([
+            SpectralConv1d(in_channels=width_fno, out_channels=width_fno, modes=modes) for _ in range(n_fno_blocks)
+            ])
+        self.w_convs = nn.ModuleList([
+            nn.Conv1d(in_channels=width_fno, out_channels=width_fno, kernel_size=1) for _ in range(n_fno_blocks)
+            ])
+        self.fc_out = nn.Linear(width_fno, dim_embedding)
+
+        if fusion_self_attention:
+            self.fusion = FusionSelfAttention(embed_dim=width_fno, hidden_dim=2*width_fno, activation=fsa_activation)
+        
+    def forward(self, x):
+        """"""
+
+        x = x.permute(0, 2, 1) #(B, D, S) -> (B, S, D)
+        x = self.fc0(x)
+        x = x.permute(0, 2, 1) #(B, S, D) -> (B, D, S) 
+
+        for i in range(self.n_fno_blocks):
+            x1 = self.spectral_convs[i](x)
+            x2 = self.w_convs[i](x)
+            
+            x = x1 + x2
+            x = self.activation(x)
+
+        if self.fusion_self_attention:
+            x = x.permute(0, 2, 1)
+            x = self.fusion(x)  
+        else:
+            x = x.mean(dim=-1)
+
+        x = self.fc_out(x)
+        
+        return x
+    
+
+class FusionSelfAttention(nn.Module):
+    def __init__(self, embed_dim: int = 64, hidden_dim: int = 64, activation: str = 'gelu'):
+        super().__init__()
+        activation = activation_by_name(activation)()
+        self.fuser = nn.Sequential(nn.Linear(embed_dim, hidden_dim), 
+                                   activation,
+                                   nn.Linear(hidden_dim, 1, bias=False))
+        
+    def forward(self, 
+                c: torch.Tensor,  # (batch_size x seq_len x embed_dim)
+                mask: torch.Tensor = None, # (batch_size x seq_len)
+                ):  
+        a = self.fuser(c)
+        alpha = torch.exp(a)*mask.unsqueeze(-1) if mask is not None else torch.exp(a)
+        alpha = alpha/alpha.sum(dim=1, keepdim=True)
         return (alpha*c).sum(dim=1)  # (batch_size x embed_dim)

reflectorch/models/encoders/integral_kernel_embedding.py

@@ -0,0 +1,390 @@
+from __future__ import annotations
+from typing import Union
+
+import torch
+from torch import nn, Tensor, stack, cat
+from reflectorch.models.activations import activation_by_name
+import reflectorch
+
+###embedding network adapted from the PANPE repository
+
+__all__ = [
+    "IntegralConvEmbedding",
+]
+
+class IntegralConvEmbedding(nn.Module):
+    def __init__(
+        self,
+        z_num: Union[int, tuple[int, ...]],
+        z_range: tuple[float, float] = None,
+        in_dim: int = 2,
+        kernel_coef: int = 16,
+        dim_embedding: int = 256,
+        conv_dims: tuple[int, ...] = (32, 64, 128),
+        num_blocks: int = 4,
+        use_batch_norm: bool = False,
+        use_layer_norm: bool = True,
+        use_fft: bool = False,
+        activation: str = "gelu",
+        conv_activation: str = "lrelu",
+        resnet_activation: str = "relu",
+    ) -> None:
+        super().__init__()
+
+        if isinstance(z_num, int):
+            z_num = (z_num,)
+        num_kernel = len(z_num)
+
+        if z_range is not None:
+            zs = [(z_range[0], z_range[1], nz) for nz in z_num]
+        else:
+            zs = z_num
+
+        self.in_dim = in_dim
+
+        self.kernels = nn.ModuleList(
+            [
+                IntegralKernelBlock(
+                    z,
+                    in_dim,
+                    kernel_coef=kernel_coef,
+                    latent_dim=dim_embedding,
+                    conv_dims=conv_dims,
+                    use_fft=use_fft,
+                    activation=activation,
+                    conv_activation=conv_activation,
+                )
+                for z in zs
+            ]
+        )
+
+        self.fc = reflectorch.models.networks.residual_net.ResidualMLP(
+            dim_in=dim_embedding * num_kernel, 
+            dim_out=dim_embedding,
+            layer_width=2 * dim_embedding,
+            num_blocks=num_blocks,
+            use_batch_norm=use_batch_norm,
+            use_layer_norm=use_layer_norm,
+            activation=resnet_activation,
+        )
+
+    def forward(self, q, y, drop_mask=None) -> Tensor:
+        x = cat([kernel(q, y, drop_mask=drop_mask) for kernel in self.kernels], dim=-1)
+        x = self.fc(x)
+        
+        return x
+
+
+class IntegralKernelBlock(nn.Module):
+    """
+    Examples:
+        >>> x = torch.rand(2, 100)
+        >>> y = torch.rand(2, 100, 3)
+        >>> block = IntegralKernelBlock((0, 1, 10), in_dim=3,  latent_dim=32)
+        >>> output = block(x, y)
+        >>> output.shape
+        torch.Size([2, 32])
+
+        >>> block = IntegralKernelBlock(10, in_dim=3,  latent_dim=32)
+        >>> output = block(x, y)
+        >>> output.shape
+        torch.Size([2, 32])
+    """
+
+    def __init__(
+        self,
+        z: tuple[float, float, int] or int,
+        in_dim: int,
+        kernel_coef: int = 2,
+        latent_dim: int = 32,
+        conv_dims: tuple[int, ...] = (32, 64, 128),
+        use_fft: bool = False,
+        activation: str = "gelu",
+        conv_activation: str = "lrelu",
+    ):
+        super().__init__()
+
+        if isinstance(z, int):
+            z_num = z
+            kernel = FullIntegralKernel(z_num, in_dim=in_dim, kernel_coef=kernel_coef)
+        else:
+            kernel = FastIntegralKernel(
+                z, in_dim=in_dim, kernel_coef=kernel_coef, activation=activation
+            )
+            z_num = z[-1]
+
+        assert z_num % 2 == 0, "z_num should be even"
+
+        self.kernel = kernel
+        self.z_num = z_num
+        self.in_dim = in_dim
+        self.latent_dim = latent_dim
+        self.use_fft = use_fft
+
+        self.fc_in_dim = self.latent_dim + self.in_dim * self.z_num
+        if self.use_fft:
+            self.fc_in_dim += self.in_dim * 2 + self.in_dim * self.z_num
+
+        self.conv = reflectorch.models.encoders.conv_encoder.ConvEncoder(
+            dim_avpool=8,
+            hidden_channels=conv_dims,
+            in_channels=in_dim,
+            dim_embedding=latent_dim,
+            activation=conv_activation,
+        )
+        self.fc = FCBlock(
+            in_dim=self.fc_in_dim, hid_dim=self.latent_dim * 2, out_dim=self.latent_dim
+        )
+
+    def forward(self, x: Tensor, y: Tensor, drop_mask: Tensor = None) -> Tensor:
+        x = self.kernel(x, y, drop_mask=drop_mask)
+
+        assert x.shape == (x.shape[0], self.in_dim, self.z_num)
+
+        xc = self.conv(x)  # (batch, latent_dim)
+
+        assert xc.shape == (x.shape[0], self.latent_dim)
+
+        if self.use_fft:
+            fft_x = torch.fft.rfft(x, dim=-1, norm="ortho")  # (batch, in_dim, z_num)
+
+            fft_x = torch.cat(
+                [fft_x.real, fft_x.imag], -1
+            )  # (batch, in_dim, 2 * z_num)
+
+            assert fft_x.shape == (x.shape[0], x.shape[1], self.z_num + 2)
+
+            fft_x = fft_x.flatten(1)  # (batch, in_dim * (z_num + 2))
+
+            x = torch.cat(
+                [x.flatten(1), fft_x, xc], -1
+            )  # (batch, in_dim * z_num * 3 + latent_dim)
+        else:
+            x = torch.cat([x.flatten(1), xc], -1)
+
+        assert (
+            x.shape[1] == self.fc_in_dim
+        ), f"Expected dim {self.fc_in_dim}, got {x.shape[1]}"
+
+        x = self.fc(x)  # (batch, latent_dim)
+
+        return x
+
+
+class FastIntegralKernel(nn.Module):
+    def __init__(
+        self,
+        z: tuple[float, float, int],
+        kernel_coef: int = 16,
+        in_dim: int = 1,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+
+        z = torch.linspace(*z)
+
+        self.kernel = FCBlock(
+            in_dim + 2, kernel_coef * in_dim, in_dim, activation=activation
+        )
+
+        self.register_buffer("z", z)
+
+    def _get_z(self, x: Tensor):
+        # x.shape == (batch_size, num_x)
+        dz = self.z[1] - self.z[0]
+        indices = torch.ceil((x - self.z[0] - dz / 2) / dz).to(torch.int64)
+
+        z = torch.index_select(self.z, 0, indices.flatten()).view(*x.shape)
+
+        return z, indices
+
+    def forward(self, x: Tensor, y: Tensor, drop_mask=None):
+        z, indices = self._get_z(x)
+        xz = torch.stack([x, z], -1)
+        kernel_input = torch.cat([xz, y], -1)
+        output = self.kernel(kernel_input)  # (batch, x_num, in_dim)
+
+        output = compute_means(
+            output * y, indices, self.z.shape[-1], drop_mask=drop_mask
+        )  # (batch, z_num, in_dim)
+
+        output = output.swapaxes(1, 2)  # (batch, in_dim, z_num)
+
+        return output
+
+
+class FullIntegralKernel(nn.Module):
+    def __init__(
+        self,
+        z_num: int,
+        kernel_coef: int = 1,
+        in_dim: int = 1,
+    ):
+        super().__init__()
+
+        self.z_num = z_num
+        self.in_dim = in_dim
+
+        self.kernel = nn.Sequential(
+            nn.Linear(in_dim + 1, z_num * kernel_coef),
+            nn.LayerNorm(z_num * kernel_coef),
+            nn.ReLU(),
+            nn.Linear(z_num * kernel_coef, z_num * in_dim),
+        )
+
+    def forward(self, x: Tensor, y: Tensor, drop_mask=None):
+        # x.shape == (batch_size, num_x)
+        # y.shape == (batch_size, num_x, in_dim)
+        # drop_mask.shape == (batch_size, num_x)
+
+        batch_size, num_x = x.shape
+
+        kernel_input = torch.cat([x.unsqueeze(-1), y], -1)  # (batch, x_num, in_dim + 1)
+        x = self.kernel(kernel_input)  # (batch, x_num, z_num * in_dim)
+        x = x.reshape(
+            *x.shape[:-1], self.z_num, self.in_dim
+        )  # (batch, x_num, z_num, in_dim)
+        # permute to get (batch, z_num, x_num, in_dim)
+        x = x.permute(0, 2, 1, 3)
+
+        y = y.unsqueeze(1)  # (batch, 1, x_num, in_dim)
+
+        assert x.shape == (
+            batch_size,
+            self.z_num,
+            num_x,
+            self.in_dim,
+        )  # (batch, z_num, in_dim, x_num)
+        assert y.shape == (
+            batch_size,
+            1,
+            num_x,
+            self.in_dim,
+        )  # (batch, 1, x_num, in_dim)
+
+        if drop_mask is not None:
+            x = x * y
+            x = x.permute(0, 2, 1, 3)  # (batch, x_num, z_num, in_dim)
+            x = masked_mean(x, drop_mask)
+        else:
+            x = (x * y).mean(-2)  # (batch, z_num, in_dim)
+
+        assert x.shape == (batch_size, self.z_num, self.in_dim), f"{x.shape}"
+
+        x = x.swapaxes(1, 2)  # (batch, in_dim, z_num)
+
+        return x
+
+
+class FCBlock(nn.Module):
+    def __init__(
+        self,
+        in_dim: int = 2,
+        hid_dim: int = 16,
+        out_dim: int = 16,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+
+        self.fc1 = nn.Linear(in_dim, hid_dim)
+        self.layer_norm = nn.LayerNorm(hid_dim)
+        self.activation = activation_by_name(activation)()
+        self.fc2 = nn.Linear(hid_dim, out_dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.layer_norm(x)
+        x = self.activation(x)
+        x = self.fc2(x)
+        return x
+        # return self.kernel(x)
+
+
+def compute_means(x, indices, z: int, drop_mask: Tensor = None):
+    """
+    Compute the mean values of tensor 'x' for each unique index in 'indices' across each batch.
+
+    This function calculates the mean of elements in 'x' that correspond to each unique index in 'indices'.
+    The computation is performed for each batch separately, and the function is optimized to avoid Python loops
+    by using advanced PyTorch operations.
+
+    Parameters:
+    x (torch.Tensor): A tensor of shape (batch_size, n, d) containing the values to be averaged.
+                      'x' should be a floating-point tensor.
+    indices (torch.Tensor): An integer tensor of shape (batch_size, n) containing the indices.
+                            The values in 'indices' should be in the range [0, z-1].
+    z (int): The number of unique indices. This determines the second dimension of the output tensor.
+    drop_mask (torch.Tensor): A boolean tensor of shape (batch_size, n) containing a mask for the indices to drop.
+                              If None, all indices are used.
+
+    Returns:
+    torch.Tensor: A tensor of shape (batch_size, z, d) containing the mean values for each index in each batch.
+                  If an index does not appear in a batch, its corresponding mean values are zeros.
+
+    Example:
+    >>> batch_size, n, d, z = 3, 4, 5, 6
+    >>> indices = torch.randint(0, z, (batch_size, n))
+    >>> x = torch.randn(batch_size, n, d)
+    >>> y = compute_means(x, indices, z)
+    >>> print(y.shape)
+    torch.Size([3, 6, 5])
+    """
+
+    batch_size, n, d = x.shape
+    device = x.device
+
+    drop = drop_mask is not None
+
+    # Initialize tensors to hold sums and counts
+    sums = torch.zeros(batch_size, z + int(drop), d, device=device)
+    counts = torch.zeros(batch_size, z + int(drop), device=device)
+
+    if drop_mask is not None:
+        # Set the values of the indices to drop to z
+        indices = indices.masked_fill(~drop_mask, z)
+
+    indices_expanded = indices.unsqueeze(-1).expand_as(x)
+    sums.scatter_add_(1, indices_expanded, x)
+    counts.scatter_add_(1, indices, torch.ones_like(indices, dtype=x.dtype))
+
+    if drop:
+        # Remove the z values from the sums and counts
+        sums = sums[:, :-1]
+        counts = counts[:, :-1]
+
+    # Compute the mean and handle division by zero
+    mean = sums / counts.unsqueeze(-1).clamp(min=1)
+
+    return mean
+
+
+def masked_mean(x, mask):
+    """
+    Computes the mean of tensor x along the x_size dimension,
+    while masking out elements where the corresponding value in the mask is False.
+
+    Args:
+    x (torch.Tensor): A tensor of shape (batch, x_size, z, d).
+    mask (torch.Tensor): A boolean mask of shape (batch, x_size).
+
+    Returns:
+    torch.Tensor: The result tensor of shape (batch, z, d) after applying the mask and computing the mean.
+    """
+    if not mask.dtype == torch.bool:
+        raise TypeError("Mask must be a boolean tensor.")
+
+    # Ensure the mask is broadcastable to the shape of x
+    mask = mask.unsqueeze(-1).unsqueeze(-1)
+    masked_x = x * mask
+
+    # Compute the sum and the count of valid (unmasked) elements along the x_size dimension
+    sum_x = masked_x.sum(dim=1)
+    count_x = mask.sum(dim=1)
+
+    # Avoid division by zero
+    count_x[count_x == 0] = 1
+
+    # Compute the mean
+    mean_x = sum_x / count_x
+
+    return mean_x

reflectorch/models/networks/__init__.py

@@ -1,14 +1,14 @@
-from reflectorch.models.networks.mlp_networks import (
-    NetworkWithPriors,
-    NetworkWithPriorsConvEmb,
-    NetworkWithPriorsFnoEmb,
-)
-from reflectorch.models.networks.residual_net import ResidualMLP
-
-
-__all__ = [
-    "ResidualMLP",
-    "NetworkWithPriors",
-    "NetworkWithPriorsConvEmb",
-    "NetworkWithPriorsFnoEmb",
-]
+from reflectorch.models.networks.mlp_networks import (
+    NetworkWithPriors,
+    NetworkWithPriorsConvEmb,
+    NetworkWithPriorsFnoEmb,
+)
+from reflectorch.models.networks.residual_net import ResidualMLP
+
+
+__all__ = [
+    "ResidualMLP",
+    "NetworkWithPriors",
+    "NetworkWithPriorsConvEmb",
+    "NetworkWithPriorsFnoEmb",
+]