Fixed params? How could make this more elegant?

[marcosrdac]

Hello! I'm trying to use Flax in my research. I wanted to have a module or flass to calculate the laplacian of a matrix using the spectral method. For this, I need my module to save a filter (phi) as a parameter that is fixed, if it makes any sense, which is generated based on the dimensions of the input matrix A. I've failed on doing this for my entire day.

The central problems are:

• I cannot let this parameter be optimized (I don't know if I should call it a parameter also);
• The parameter must be saved within the module for the next applications, like when you define a Dense layer from within the call method;
• The solution should be efficient.

Here was my attempt, where calculate phi at each iteration (not wanted!).

def applyfilter(phi, A):
    return jnp.fft.ifft2(phi * jnp.fft.fft2(A)).real


class SpectralLaplacian(nn.Module):
    dz: float
    dx: float
    
    def gen_filter(self, A, dz=1., dx=1.):
        nz, nx = A.shape
        dkz = 2*np.pi/(nz*dz)
        dkx = 2*np.pi/(nx*dx)
        kz, kx = jnp.mgrid[-nz/2:nz-nz/2, -nx/2:nx-nx/2]
        kz = jnp.fft.ifftshift(kz) * dkz
        kx = jnp.fft.ifftshift(kx) * dkx
        phi = -(kz**2 + kx**2)
        return phi

    def setup(self):
        # I wanted to maybe set phi up here somehow...
        pass

    def __call__(self, A):
        phi = self.gen_filter(A, self.dz, self.dx)  # phi should only be generated once
        return applyfilter(phi, A)

How could I solve this problem? What would be an elegant way to do it? I'm open for tips!

Thank you beforehand!

[jheek]

Consider using a variable that is not a param for phi (https://flax.readthedocs.io/en/latest/flax.linen.html#flax.linen.Module.variable). You can define your own collection name for this filter like

phi = self.variable("filter", "phi", ...)
phi.value = ...

[marcosrdac]

Thank you for your attention, @jheek! I did not really understand the concept of variable in my first readings... Now it makes sense. The code became:

class SpectralLaplacian(nn.Module):
    dz: float
    dx: float

    def gen_filter(self, shape):
        nz, nx = shape
        dz, dx = self.dz, self.dx
        dkz = 2*np.pi/(nz*dz)
        dkx = 2*np.pi/(nx*dx)
        kz, kx = np.mgrid[-nz/2:nz-nz/2,
                          -nx/2:nx-nx/2]
        kz = jnp.fft.ifftshift(kz) * dkz
        kx = jnp.fft.ifftshift(kx) * dkx
        phi = -(kz**2 + kx**2)
        return phi

    @nn.compact
    def __call__(self, A):
        phi = self.variable('filters', 'laplacian', self.gen_filter, A.shape)
        return applyfilter(phi.value, A)