Module that is flexible in which submodule is used

[ingmarschuster]

Hi,

related to #665 and #686 I'm currently porting JaxRk to use FLAX. What I want is to write Modules (like AverageSimilarity in the following example) that can be flexibly parameterized to use different types of submodules (like LaplaceSimilarity and CosineSimilarity in the following), so long as they have the same call signature. I came up with a similar pattern to what was suggested in the end of thread #665.
Here is the minimal working example:

import functools
from typing import Any, Callable, Sequence, Optional
import numpy as np
import jax
from jax import lax, random, numpy as jnp
import flax
from flax.core import freeze, unfreeze
from flax import linen as nn

class CosineSimilarity(nn.Module):
  def setup(self):
    pass

  def __call__(self, x, y):
    return np.dot(x,y)

class LaplaceSimilarity(nn.Module):
  scale:float = 1.
  def setup(self):
    pass

  def __call__(self, x, y):
    diff = x - y
    return jnp.exp(-jnp.sqrt(jnp.sum(diff**2))/self.scale)

class AverageSimilarity(nn.Module):
  sim_factory:Callable

  def setup(self):
    self.sim = self.sim_factory()

  def __call__(self, inputs1, inputs2):
    return self.sim(inputs1, inputs2).mean()

key1, key2, key3 = random.split(random.PRNGKey(0), 3)
x, y = random.uniform(key1, (4,4)), random.uniform(key2, (4,4))

for sim_factory, name  in [(functools.partial(LaplaceSimilarity, 3.), "laplace"),
                           (functools.partial(CosineSimilarity), "cosine")]:
    model = AverageSimilarity(sim_factory = sim_factory)
    init_variables = model.init(key3, x, y)
    res = model.apply(init_variables, x, y)

    print(name)
    print('output:\n', res)

Is this the best way?
I guess this itch is not special to my use case.

[avital]

I think that after 9bf7161 was committed, you should be able to simply pass an unbound module instance from the top-level without wrapping in a lambda. Could you give it a shot?

e.g. model = AverageSimilarity(LaplaceSimilarity(3.))

(On another note, there's no need to define empty setup methods)

(Oh, and your LaplaceSimilarity class doesn't really need to be a module as it doesn't hold variables. You could simply make it a vanilla Python dataclass)

[ingmarschuster]

Indeed, it works now after using the current master branch. Adapted example:

from jax import random, numpy as jnp
from flax import linen as nn

class CosineSimilarity(nn.Module):
  def __call__(self, x, y):
    return jnp.dot(x,y)

class LaplaceSimilarity(nn.Module):
  scale:float = 1.

  def __call__(self, x, y):
    diff = x - y
    return jnp.exp(-jnp.sqrt(jnp.sum(diff**2))/self.scale)

class AverageSimilarity(nn.Module):
  sim:nn.Module

  def __call__(self, inputs1, inputs2):
    return self.sim(inputs1, inputs2).mean()

key1, key2, key3 = random.split(random.PRNGKey(0), 3)
x, y = random.uniform(key1, (4,4)), random.uniform(key2, (4,4))

for sim_factory, name  in [(LaplaceSimilarity(3.), "laplace"),
                           (CosineSimilarity(), "cosine")]:
    model = AverageSimilarity(sim = sim_factory)
    init_variables = model.init(key3, x, y)
    res = model.apply(init_variables, x, y)

    print(name)
    print('output:\n', res)