Multiple Inheritance -> doesn't recognize as Module throws ValueError: parent must be None, Module or Scope

[SauravMaheshkar]

I'm working on a Flax implementation for ProteinBERT: A universal deep-learning model of protein sequence and function. My work so far is in SauravMaheshkar/ProteinBERT.

I've made a simple test.py to check instantiation using the .init() function. My test script is as follows :

from proteinbert import ProteinBERT
import jax
from jax import random


def test():

    seq = jax.random.randint(
        key=random.PRNGKey(0), minval=0, maxval=21, shape=(2, 2048)
    )
    annotation = jax.random.randint(
        key=random.PRNGKey(0), minval=0, maxval=1, shape=(2, 8943)
    )

    init_rngs = {"params": random.PRNGKey(0), "layers": random.PRNGKey(1)}

    ProteinBERT().init(init_rngs, seq, annotation)


if __name__ == "__main__":
    test()

And I've been getting this error message

Error Message

Traceback (most recent call last):
  File "/Users/sauravmaheshkar/github/protein_bert/test.py", line 21, in <module>
    test()
  File "/Users/sauravmaheshkar/github/protein_bert/test.py", line 17, in test
    ProteinBERT().init(init_rngs, seq, annotation)
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 1000, in init
    method=method, mutable=mutable, **kwargs)
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 969, in init_with_output
    {}, *args, rngs=rngs, method=method, mutable=mutable, **kwargs)
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 939, in apply
    )(variables, *args, **kwargs, rngs=rngs)
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/core/scope.py", line 687, in wrapper
    y = fn(root, *args, **kwargs)
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 1178, in scope_fn
    return fn(module.clone(parent=scope), *args, **kwargs)
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 266, in wrapped_module_method
    self._try_setup()
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 679, in _try_setup
    self.setup()
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 275, in wrapped_module_method
    y = fun(self, *args, **kwargs)
  File "/Users/sauravmaheshkar/github/protein_bert/proteinbert/model.py", line 82, in setup
    Reduce("b n d -> b d", "mean"),
  File "<string>", line 5, in __init__
  File "/Users/sauravmaheshkar/opt/anaconda3/envs/proteinbert/lib/python3.7/site-packages/flax/linen/module.py", line 599, in __post_init__
    raise ValueError("parent must be None, Module or Scope")
ValueError: parent must be None, Module or Scope

The problem lies in the Reduce defined in proteinbert/utils.py class which is defined as follows:

class Reduce(ReduceMixin, nn.Module):
    """
    Flax Module to act as a Reduce layer (from einops)
    """
    def __call__(self, input):
        return self._apply_recipe(input)

The idea is to create a Reduce layer/Module for flax which performs the reduce operation from einops. Although the module inherits from flax.linen.Module it still throws a ValueError.

Any help would be much appreciated 😊.

[avital]

@marcvanzee will be looking into this.

[marcvanzee]

Hi @SauravMaheshkar, Flax Modules are special kinds of dataclasses in the sense that they will attach parent and name attributes automatically. The parent attribute is usually assigned automatically: it is set to the parent Module, or set to None if this is the top-level module. However, if you provide Modules with a string argument, the parent attribute will be set to the string value. Here is a minimal example:

class Bar(nn.Module):
  @nn.compact
  def __call__(self, x):
    pass

>>> Bar('test')
...
ValueError: parent must be None, Module or Scope

In your case, in ProteinBert.setup(), you initialize Reduce as follows: Reduce("b n d -> b d", "mean"). You intend to call the constructor of ReduceMixin with these two string arguments, but Module will take the first argument and treat it as its parent.

(Note that parent is attached manually, so the Method Resolution Order doesn't matter here.)

I think this is something we should try to improve from the Flax side as well, but in the meantime an easy solution is to make ReduceMixin an attribute to the Reduce module. I actually think this is conceptually a bit easier to understand as well because it avoids mental overhead in understanding how the double inheritance initialization works.

class Reduce(nn.Module):
  pattern: str
  reduction: str

  def setup(self):
    self.reducer = ReduceMixin(self.pattern, self.reduction)

  def __call__(self, input):
    return self.reducer._apply_recipe(input)

class Foo(nn.Module):

  def setup(self):
    self.reducer = Reduce(pattern="b n d -> b d", reduction="mean")

  def __call__(self, x):
    ...


seq = jax.random.randint(key=random.PRNGKey(0), minval=0, maxval=21, shape=(2, 2048))
annotation = jax.random.randint(key=random.PRNGKey(0), minval=0, maxval=1, shape=(2, 8943))

init_rngs = {"params": random.PRNGKey(0), "layers": random.PRNGKey(1)}

Foo().init(init_rngs, seq)

Please let me know if this works for you!

[SauravMaheshkar]

Thanks a ton @marcvanzee. The Reduce error is fixed now. Now on to the next error 😅. Closing the discussion.