Repository: google/flax

Issue:
Error when JITting `Model.__call__`
eg

```python
import jax
from flax import nn
layer=nn.Dense.partial(features=1)
key=jax.random.PRNGKey(0)
x=jax.random.normal(key, (20, 2))
_,params=layer.init(key, x)
layer_m=nn.Model(layer, params)
jax.jit(layer_m)(x)
```

errors with 

```
TypeError                                 Traceback (most recent call last)
<ipython-input-2-2e4e0581e3f5> in <module>
      6 _,params=layer.init(key, x[0,...])
      7 layer_m=nn.Model(layer, params)
----> 8 jax.jit(layer_m)(x)

~/opt/anaconda3/lib/python3.7/site-packages/jax/api.py in f_jitted(*args, **kwargs)
    148     flat_fun, out_tree = flatten_fun(f, in_tree)
    149     out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend,
--> 150                        name=flat_fun.__name__)
    151     return tree_unflatten(out_tree(), out)
    152

~/opt/anaconda3/lib/python3.7/site-packages/jax/linear_util.py in __name__(self)
    121   @property
    122   def __name__(self):
--> 123     return getattr(self.f, '__name__', '<unnamed wrapped function>')
    124
    125   def wrap(self, gen, gen_static_args, out_store) -> 'WrappedFun':

~/opt/anaconda3/lib/python3.7/site-packages/flax/nn/base.py in __getattr__(self, name)
    897   def __getattr__(self, name):
    898     value = getattr(self.module, name)
--> 899     if issubclass(value, Module):
    900       def wrapper(*args, **kwargs):
    901         return value.call(self.params, *args, **kwargs)

~/opt/anaconda3/lib/python3.7/abc.py in __subclasscheck__(cls, subclass)
    141         def __subclasscheck__(cls, subclass):
    142             """Override for issubclass(subclass, cls)."""
--> 143             return _abc_subclasscheck(cls, subclass)
    144
    145         def _dump_registry(cls, file=None):

TypeError: issubclass() arg 1 must be a class
```


Discussion:
Sorry, it took me a bit to figure out what was going on.
A Model should be pmap'able - what's happening here is a bit of a subtle bug:

First, a short-term "fix" is just wrapping it in a lambda passthrough:
```python
import jax
from flax import nn
layer=nn.Dense.partial(features=1)
key=jax.random.PRNGKey(0)
x=jax.random.normal(key, (4, 20, 2))
_,params=layer.init(key, x[0,...])
layer_m=nn.Model(layer, params)
jax.pmap(lambda z: layer_m(z))(x)
```

Now, what's going on:
 - in a great change https://github.com/google/jax/pull/2073 made ~2 months ago to improve XLA call stack metadata JAX tries to get the `__name__` attribute from the pmap'd function, which in this case is our callable Model instance.   
 - the problem is that in another refactoring of the base flax code a month ago https://github.com/google/flax/commit/baf43e73cb0088a607c4da26be981a83bfaf6a52 we override `__getattr__` on Model to passthrough and grab the requested attr from Module, but inside that we are trying to eval `issubclass(fetched_attr, flax.nn.Module)` and `issubclass(<string object>, flax.nn.Module)` throws an error in python since it's nonsense.

We almost always use a Model inside an optimizer or indirectly in another function, and I think we must not have a unit test of a direct jit/pmap on a Model - my apologies for letting this slip through, we'll try to get a fix in asap.

Resolving patch:
diff --git a/flax/nn/base.py b/flax/nn/base.py
--- a/flax/nn/base.py
+++ b/flax/nn/base.py
@@ -919,7 +919,7 @@ def truncate_at(self, module_path):
 
   def __getattr__(self, name):
     value = getattr(self.module, name)
-    if issubclass(value, Module):
+    if inspect.isclass(value) and issubclass(value, Module):
       def wrapper(*args, **kwargs):
         return value.call(self.params, *args, **kwargs)
       return wrapper


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -102,6 +102,16 @@ def test_init_by_shape_module(self):
     self.assertEqual(y2, jnp.array([2.]))
     self.assertEqual(params, {'bias': jnp.array([1.])})
 
+  def test_model(self):
+    rng = random.PRNGKey(0)
+    x = jnp.array([1.])
+    _, params = DummyModule.init(rng, x)
+    model = nn.Model(DummyModule, params)
+    y = model(x)
+    self.assertEqual(y, jnp.array([2.]))
+    y2 = jax.jit(model)(x)
+    self.assertEqual(y2, jnp.array([2.]))
+
   def test_shared_module(self):
     rng = random.PRNGKey(0)
     x = jnp.array([1.])
@@ -272,6 +282,11 @@ def apply(self, x):
         MultiMethod.__qualname__ + '.l2')
 
     x = jnp.array([1., 2.])
+
+    _, params = MultiMethod.init(random.PRNGKey(0), x)
+    model = nn.Model(MultiMethod, params)
+    self.assertEqual(model.l2(), 2.)
+    
     y, _ = MultiMethodModel.init(random.PRNGKey(0), x)
     self.assertEqual(y, 2.)