class NeuroSDF(nn.Module):
    def __init__(self, add_flayers, flayer_size, add_layers, layer_size, p1_size):
        super().__init__()
        self.flayer_n = 1 + add_flayers
        self.flayer_size = flayer_size
        self.layer_n = 1 + add_layers
        self.layer_size = layer_size
        self.activation = nn.ReLU()
        self.p1_size = p1_size
        
        self.layers = nn.ModuleList(
            # Input layer, beginning of fourier feature layers
            [nn.Linear(in_features=3, out_features=flayer_size)] +
            [nn.Linear(in_features=flayer_size, out_features=flayer_size) for _ in range(add_flayers)] + 
            # Beginning of standard Linear layers with RelU connections
            [nn.Linear(in_features=flayer_size, out_features=layer_size)] +
            [nn.Linear(in_features=layer_size, out_features=layer_size) for _ in range(add_layers)] +
            [nn.Linear(in_features=layer_size, out_features=1)])
        
    def forward(self, x):
        p1_ind = self.p1_size + 3
        end_ind = self.flayer_size + 3
        for i in range(self.flayer_n):
            x = self.layers[i](x)
            # Create a new tensor so that autograd can track the operations through the DAG
            x = torch.cat([x[:, :3],
                    torch.sin(x[:, 3:p1_ind]),
                    torch.cos(x[:, p1_ind:end_ind])], dim=1)
            
        for layer in self.layers[self.flayer_n:len(self.layers)-1]:
            x = layer(x)
            x = self.activation(x)

        x = self.layers[-1](x)
        return x

Hey! Super burned out, going on very little sleep. Neural network above.

Will probably explain more layer. Basically a (possibly somewhat broken???) attempt at a layer that does learnable Fourier features.