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

^{Originally posted by **SauravMaheshkar** June 26, 2021}
I'm working on a Flax implementation for [ProteinBERT: A universal deep-learning model of protein sequence and function](https://www.biorxiv.org/content/10.1101/2021.05.24.445464v1). My work so far is in [SauravMaheshkar/ProteinBERT](https://github.com/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
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 "", 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](https://github.com/SauravMaheshkar/ProteinBERT/blob/jaxpackage/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 😊.