Hypernetworks with flax?

[ingmarschuster]

Hi,

how would I implement a network that spits out another network after some calculation? Consider the following minimal example:

import flax.linen as ln
import jax.numpy as np
from jax import random

class Function(ln.Module):
    f:float = 1.
    def setup(self):
        self.pow = np.exp(self.param("pow", random.uniform))

    def __call__(self, x:float):
        return self.f * np.power(x, self.pow)

class FunctionSpaceOperator(ln.Module):
    left:Function
    right:Function
    factor:float
    def setup(self):
        pass

    def __call__(self, x:float):
        return Function(self.right(x) * self.factor * self.left.f)

    def easy_init(self, ):
        self.__call__(4.)

f = FunctionSpaceOperator(Function(3.), Function(2.), 4.)
f = f.bind(f.init(random.PRNGKey(1), method=f.easy_init))
print(f(4))

which throws a

AssignSubModuleError: Submodule Function must be defined in `setup()` or in a method wrapped in `@compact

Thanks for your input!

[jheek]

In this case you should add @ln.compact to FunctionSpaceOperator.__call__.
The decorator makes inline sub modules work.

[ingmarschuster]

Ha, this is true of course, maybe my minimal example was too minimal. Actually, there is more than one method which would need the inline submodule, so your suggested solution unfortunately doesn't work. Updated example (bit contrived now):

import flax.linen as ln
import jax.numpy as np
from jax import random

class Function(ln.Module):
    f:float = 1.
    def setup(self):
        self.pow = np.exp(self.param("pow", random.uniform))

    def __call__(self, x:float):
        return self.f * np.power(x, self.pow)

class FunctionSpaceOperator(ln.Module):
    left:Function
    right:Function
    factor:float

    @ln.compact
    def __call__(self, x:float):
        return Function(self.right(x) * self.factor * self.left.f)


    @ln.compact
    def __add__(self, other:"FunctionSpaceOperator"):
        return FunctionSpaceOperator(self.left, other.right, self.factor + other.factor)
    
    @ln.compact
    def easy_init(self, ):
        self.__call__(4.)

f = FunctionSpaceOperator(Function(3.), Function(2.), 4.)
f = f.bind(f.init(random.PRNGKey(1), method=f.easy_init))
print(f(4))

Concretely, I'm trying to fix up the FiniteVec.mean, FiniteVec.sum, FiniteVec.centered methods in this file and the FiniteOp.__matmul__ in this one all of which used to return FiniteVec or FiniteOp instances before. I can change the design for making the FLAX integration work, but unsure how.

[jheek]

I think the issue here comes down to who owns the submodules. In this example you should just remove the compact decorator from __add__ and easy_init.
This means that if another module calls __add__ it will have to be in a compact function and it will "own" the add module.

You can have only 1 compact method in a Module but other methods can call a compact method.

[ingmarschuster]

And the other module calling __add__ and owning the created submodule could also do so in a setup method, not only in one decorated with @compact, correct?

[ingmarschuster]

hm, this actually means one has to replicate initialization code. Consider the following from the actual library I'm working on. I could move the initialization from setup() to, e.g. __matmul__(), but that would mean that for calling inner() and T() I have to do the initialization again in the module calling these methods.

class FiniteOp(LinOp[InpVecT, OutVecT], Module):
    """Finite rank LinOp in RKHS"""
    inp_feat:InpVecT
    outp_feat:OutVecT
    matr_init:AnyOrInitFn
    normalize:bool = False

    def setup(self, ):
        self.matr = self.const_or_param("matr", self.matr_init)

    def __len__(self):
        return len(self.inp_feat) * len(self.outp_feat)

    def __matmul__(self, right_inp:CombT) -> Union[OutVecT, "FiniteOp[RhInpVectT, OutVecT]"]:
        if isinstance(right_inp, FiniteOp):
            matr = self.matr @ self.inp_feat.inner(right_inp.outp_feat) @ right_inp.matr
            rval = FiniteOp(right_inp.inp_feat, self.outp_feat, matr)
            return rval
        else:
            if isinstance(right_inp, DeviceArray):
                right_inp = FiniteVec(self.inp_feat.k, np.atleast_2d(right_inp))  
            lin_LinOp = (self.matr @ inner(self.inp_feat, right_inp)).T
            if self._normalize:
                lin_LinOp = lin_LinOp / lin_LinOp.sum(1, keepdims = True)
            r = [LinearReduce(lin_LinOp)]
            if len(right_inp) == 1:
                r.append(Sum())
            rval = self.outp_feat.extend_reduce(r)
            return rval
    
    def inner(self, Y:"FiniteOp[InpVecT, OutVecT]"=None, full=True):
        assert NotImplementedError("This implementation as to be tested")
        if Y is None:
            Y = self
        G_i = self.inp_feat.inner(Y.inp_feat)
        G_o = self.outp_feat.inner(Y.outp_feat)

        # check the following expression again
        if Y.matr.size < self.matr.size:
            return np.sum((G_o.T @ self.matr @ G_i) * Y.matr)
        else:
            return np.sum((G_o @ Y.matr @ G_i.T) * self.matr)

    def reduce_gram(self, gram, axis = 0):
        return gram
    
    @property
    def T(self) -> "FiniteOp[OutVecT, InpVecT]":
        return FiniteOp(self.outp_feat, self.inp_feat, self.matr.T, self.normalize)

[jheek]

And the other module calling add and owning the created submodule could also do so in a setup method, not only in one decorated with @compact, correct?

yes

As for your example I find it hard to grasp exactly what you are trying to do. For example the transpose property creates a new Module that optionally creates a new "matr" paramater. That seems strange. I'm not sure why you want the creation of parameters and compute to be linked like that.

[ingmarschuster]

Thanks for being so persistent in helping me out here. I tried your suggestion by calling __matmul__ from another modules setup. This still results in

~/Devel/python/JaxRK/jaxrk/rkhs/operator.py in __matmul__(op, right_inp)
     64             if op.normalize:
     65                 lin_LinOp = lin_LinOp / lin_LinOp.sum(1, keepdims = True)
---> 66             r = [LinearReduce(lin_LinOp)]
     67             if len(right_inp) == 1:
     68                 r.append(Sum())

~/Devel/python/JaxRK/jaxrk/reduce/lincomb.py in __init__(self, linear_map, parent, name)

/opt/anaconda3/lib/python3.8/site-packages/flax/linen/module.py in __post_init__(self)
    564         return
    565       if not self.parent._initialization_allowed:
--> 566         raise errors.AssignSubModuleError(self.__class__.__name__)
    567       # Autonaming of submodules.
    568       if self.name is None:  # pytype: disable=attribute-error

AssignSubModuleError: Submodule LinearReduce must be defined in `setup()` or in a method wrapped in `@compact` (https://flax.readthedocs.io/en/improve-error/flax.errors.html#flax.errors.AssignSubModuleError)

However, by simply pulling __matmul__ out of the class definition and making it a function (def mm(op:FiniteOp[InpVecT, OutVecT], right_inp:CombT) …) this seems to be working. Of course like this one can't overload Python operators as I'd like to have done. Maybe a method property linen.factory could be handy here for methods like __matmul__, __add__, etc. that would make clear that the ownership of modules created in these functions has to go to the caller?

Regarding the T() method: a fixed matrix is passed to FiniteOp, which leads to self.const_or_param not creating a new parameter, but rather using the passed array directly. The overall principle in this library is to make it possible to use gradient based optimization when needed, but also enable closed form solutions, which are often available in RKHS methods.

[ingmarschuster]

I found that my previous design for integrating FLAX into JaxRK does not make sense and have changed this. Thus I'm marking the question as answered.