internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: google/flax

Issue:
`typing.get_type_hints()` is broken for linen modules
I have some serialization code that involves a recursive call to `get_type_hints()`, which breaks for flax modules:
```python
from typing import get_type_hints

from flax import linen as nn


class Network(nn.Module):
    layers: int

# Fails!
# NameError: name 'Module' is not defined
print(get_type_hints(Network))
```

The reason for this seems to be that forward references are (seemingly unnecessarily) used when fields are being dynamically added to the module dataclass, but the typing module tries to resolve these names in the wrong local namespace:
https://github.com/google/flax/blob/96c78cd1bb43dfacfb8a999f3155facec00ecb3b/flax/linen/module.py#L533-L534

This can be confirmed because adding one extra line fixes the error:
```python
from typing import get_type_hints

from flax import linen as nn
from flax.linen.module import Module, Scope, _Sentinel  # New


class Network(nn.Module):
    layers: int

# Works!
# {'layers': <class 'int'>, 'parent': typing.Union[typing.Type[flax.linen.module.Module], typing.Type[flax.core.scope.Scope], typing.Type[flax.linen.module._Sentinel], NoneType], 'name': <class 'str'>}
print(get_type_hints(Network))
```


Discussion:

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -530,8 +530,8 @@ def _customized_dataclass_transform(cls):
     """Handles final optional dataclass attributes: `parent` and `name`."""
     # Use cls.__dict__ to get annotations of cls itself (no parent class).
     annotations = dict(cls.__dict__.get('__annotations__', {}))
-    parent_annotation = Union[Type["Module"], Type["Scope"],
-                              Type["_Sentinel"], None]
+    parent_annotation = Union[Type[Module], Type[Scope],
+                              Type[_Sentinel], None]
     if ('parent' in annotations
         and annotations['parent'] != parent_annotation):
       raise errors.ReservedModuleAttributeError(annotations)


Test patch:
Repository: google/flax

Issue:
Can not assign dict whos key is not string as module attribute
Hi,
It seems current flax.linen not allow assign a dict with non-string keys to the module attribute.
See below simple example:
it will trigger error:
`AssertionError: A state dict must only have string keys.`

Questions:
1. Is it an intended behavior? Why?
2. If it is intended, is there any work around? As it is quite possible we need assign the information contained by a dict to the module, and the key of the dict may not string.

```python

import flax.linen as nn
import jax
import jax.numpy as jnp


class Foo(nn.Module):
    def setup(self):
        self.a = {(1, 2): 3} # here the dict using tuple as key

    @nn.compact
    def __call__(self, x):
        return x


foo = Foo()
rng = jax.random.PRNGKey(0)
x = jnp.ones(shape=(3, 3))
vars = foo.init({"params": rng}, x)
out = foo.apply(vars, x)
print(out)
```




Discussion:
This should be fixed
@jheek just tried this in a public Colab and installed flax from main, but the problem still seems to be there.
It's because we traverse any assignment looking for Module leaves and have overly strict requirements on the structure of the tree (e.g. string keys) for any leaf Module and that's spilling over as a constraint on any leaf type.

Resolving patch:
diff --git a/flax/serialization.py b/flax/serialization.py
--- a/flax/serialization.py
+++ b/flax/serialization.py
@@ -117,11 +117,14 @@ def _restore_list(xs, state_dict: Dict[str, Any]) -> List[Any]:
 
 
 def _dict_state_dict(xs: Dict[str, Any]) -> Dict[str, Any]:
-  return {key: to_state_dict(value) for key, value in xs.items()}
+  str_keys = set(str(k) for k in xs.keys())
+  if len(str_keys) != len(xs):
+    raise ValueError(f'Dict keys do not have a unique string representation: {str_keys}')
+  return {str(key): to_state_dict(value) for key, value in xs.items()}
 
 
 def _restore_dict(xs, states: Dict[str, Any]) -> Dict[str, Any]:
-  return {key: from_state_dict(value, states[key])
+  return {key: from_state_dict(value, states[str(key)])
           for key, value in xs.items()}
 
 


Test patch:
diff --git a/tests/linen/linen_module_test.py b/tests/linen/linen_module_test.py
--- a/tests/linen/linen_module_test.py
+++ b/tests/linen/linen_module_test.py
@@ -159,6 +159,22 @@ def __call__(self, x):
       {'lyrs1_a': {'kernel': (10, 3)},
       'lyrs1_b': {'kernel': (3, 3)}})
 
+  def test_setup_dict_nonstring_keys(self):
+    class Foo(nn.Module):
+      def setup(self):
+          self.a = {(1, 2): nn.Dense(2)} # here the dict using tuple as key
+
+      @nn.compact
+      def __call__(self, x):
+          return self.a[(1, 2)](x)
+
+    foo = Foo()
+    x = jnp.ones(shape=(1, 3))
+    params = foo.init(random.PRNGKey(0), x)['params']
+    param_shape = jax.tree_map(jnp.shape, params)
+    self.assertEqual(param_shape,
+      {'a_(1, 2)': {'kernel': (3, 2), 'bias': (2,)}})
+
   def test_setup_cloning(self):
     class MLP(nn.Module):
       def setup(self):
Repository: google/flax

Issue:
Flax actually requires jax 0.3.2
https://github.com/google/flax/blob/ef6bf4054c30271a58bfabb58f3d0049ef5d851a/flax/linen/initializers.py#L19

the constant initialiser was added in this commit https://github.com/google/jax/commit/86e8928e709ac07cc51c10e815db6284507c320e that was first included in jax 0.3.2

This came up in NetKet's automated oldest-version-dependencies testing.


Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -26,7 +26,7 @@
 
 install_requires = [
     "numpy>=1.12",
-    "jax>=0.3",
+    "jax>=0.3.2",
     "matplotlib",  # only needed for tensorboard export
     "msgpack",
     "optax",


Test patch:
Repository: google/flax

Issue:
New Sphinx Theme
The idea is to get a new and shiny theme that makes Flax's RTD page standout a little more. 

I've gathered a couple of options:

### [JAX's Theme](https://jax.readthedocs.io/en/latest/)
![Screenshot from 2022-05-25 13-16-27](https://user-images.githubusercontent.com/5862228/170337666-3fdac1a8-d6c6-45ad-805e-8c3723c487a3.png)
### [Pydata Sphinx Theme](https://pydata-sphinx-theme.readthedocs.io/en/latest/user_guide/index.html)
![Screenshot from 2022-05-25 13-17-39](https://user-images.githubusercontent.com/5862228/170337899-f9ae6ecf-5681-4a32-b5e6-4b56f4f83be6.png)
### [Furo](https://pradyunsg.me/furo/quickstart/)
![Screenshot from 2022-05-25 13-19-51](https://user-images.githubusercontent.com/5862228/170338218-ca637285-bf1b-4fc2-8637-a18b8c922ad5.png)



Discussion:
+1 to JAX
+1 to Furo
@marcvanzee Here are some samples from JAX's theme (sphinx_book_theme) and furo.

My 2 cents: I like furo a little better but they both look good.

## sphinx_book_theme
#### landing page
![Screenshot from 2022-05-26 13-11-47](https://user-images.githubusercontent.com/5862228/170550547-6617a4d2-abe0-49d7-a8c1-ab32d27290ac.png)
 #### module
![Screenshot from 2022-05-26 13-12-13](https://user-images.githubusercontent.com/5862228/170550580-65fcbf4b-1dfd-46f9-837f-37dc5a724668.png)
## furo
#### landing page
![Screenshot from 2022-05-26 13-08-42](https://user-images.githubusercontent.com/5862228/170550629-a6b2e9f5-93c7-4b29-962d-32ca8e9f97a4.png)
#### module
![Screenshot from 2022-05-26 13-09-18](https://user-images.githubusercontent.com/5862228/170550675-ee2e4dc7-2ef2-4dea-908f-be6469e5e8b7.png)


Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -84,8 +84,8 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'sphinx_rtd_theme'
-html_style = 'css/flax_theme.css'
+html_theme = 'sphinx_book_theme'
+# html_style = 'css/flax_theme.css'
 
 # The name of an image file (relative to this directory) to place at the top
 # of the sidebar.


Test patch:
Repository: google/flax

Issue:
Outdated `rich` dependency version
The version of `rich` is currently limited to `rich~=11.1`, causing problems with `pip` dependency resolution when installing with other packages.

https://github.com/google/flax/blob/cda7a4c85bbce744e412ab82e298ddf76d4770d2/setup.py#L33

Should be a trivial fix since `flax.linen.summary` doesn't seem to need any changes, I'll open a PR.



Discussion:

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,7 @@
     "matplotlib",  # only needed for tensorboard export
     "msgpack",
     "optax",
-    "rich~=11.1",
+    "rich>=11.1",
     "typing_extensions>=4.1.1",
     "PyYAML>=5.4.1",
 ]


Test patch:
Repository: google/flax

Issue:
Raise clearer Exception when calling method of unbound module
Using this minimal example
```
import jax.numpy as np
from jax.numpy import log, exp
import jax.random as rand
import flax.linen as ln

class MultipleForw(ln.Module):
    def setup(self):
        self.s1 = self.param("s1", ln.initializers.ones, (1,))

    def __call__(self, X, ):
        return X * log(1 + exp(self.s1 - 1))



mf = MultipleForw()
X = np.arange(5)
mf.init(rand.PRNGKey(0), X)
mf(X)
```

### Problem you have encountered:
The last line raised the rather opaque error message `AttributeError: 'MultipleForw' object has no attribute 's1'`

### What you expected to happen:
The raised Exception should contain a hint that makes clear that calling a linen-Module correctly is by using `mf.apply(parameters, input)`. See Discussion #1013 




Discussion:
In #1072, I tried fixing this by creating a custom error class for Module AtttributeError, but after a discussion with @avital we found that this is not a very natural solution because users expect a normal `AttributeError` when they are trying to access an unknown attribute in a Module.

Solving this issue is a bit more work, and probably not our highest priority. For now I'm lowering the priority of this issue because it seems we won't fix it soon, and we can higher it when it turns out that more users run into this problem.
Also unassigning myself since I don't plan to work on this soon.
I'd like to take this issue. A simple solution would be to customize the current error message to suggest calling `apply` if `self.scope is None`. I'll create a PR as this is a rather simple fix, if we want to tackle it a different way we can discuss there.
This already sounds an order of magnitude better than the current situation

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -775,8 +775,11 @@ def __getattr__(self, name: str) -> Any:
     if name in self.__dict__:
       return self.__dict__[name]
     else:
-      raise AttributeError(
-          f'"{self.__class__.__name__}" object has no attribute "{name}"')
+      msg = f'"{self.__class__.__name__}" object has no attribute "{name}".'
+      if self.scope is None:
+        msg += (f' If "{name}" is defined in \'.setup()\', remember these fields '
+          'are only accessible from inside \'init\' or \'apply\'.')
+      raise AttributeError(msg)
 
   def __dir__(self) -> List[str]:
     """Call setup() before listing attributes."""


Test patch:
Repository: google/flax

Issue:
Pytype attribute generics tracker
Pytype currently doesn't support Generic types class attributes:
```
class Foo:
  bar: T
```

As a workaround we use `Any` for the attribute instead. This workaround should be reverted once the functionality is implemented


Discussion:
@jheek  -- Is there a public bug we can reference tracking the current limitation in pytype?
@jheek is this still relevant?
Yes we are still using this workaround

Resolving patch:
diff --git a/flax/core/lift.py b/flax/core/lift.py
--- a/flax/core/lift.py
+++ b/flax/core/lift.py
@@ -315,13 +315,13 @@ def swap(target):
 @dataclasses.dataclass(frozen=True)
 class In(Generic[T]):
   """Specifies a variable collection should only be lifted as input."""
-  axis: Any  # pytype does not support generic variable annotation
+  axis: T
 
 
 @dataclasses.dataclass(frozen=True)
 class Out(Generic[T]):
   """Specifies a variable collection should only be lifted as output."""
-  axis: Any  # pytype does not support generic variable annotation
+  axis: T
 
 
 def _split_in_out_axes(xs: Mapping[CollectionFilter, Any]):


Test patch:
Repository: google/flax

Issue:
Improve documentation for `Dropout` and `rngs` argument in `linen.Module.apply()`

Here is an example of `Dropout` in a model definition:
https://github.com/google/flax/blob/d068512a932da3e05b822790a591bac391aeab36/examples/nlp_seq/models.py#L211

Here is the `apply()`, where `rngs` is passed in
https://github.com/google/flax/blob/d068512a932da3e05b822790a591bac391aeab36/examples/nlp_seq/train.py#L206-L207
However the `rng` is not very clearly explained in `apply()`
https://github.com/google/flax/blob/615f40be774e7ed66fd344e8291ac0d48ebcef7d/flax/linen/module.py#L749
The `rngs` seems to be passed to `flax/core/scope.py`
Here is the code for `Dropout` (linen)
https://github.com/google/flax/blob/9b4807840c5cb26ef5e29028e3558d404aee00a0/flax/linen/stochastic.py#L56-L57
Here is the code for  `make_rng()`
https://github.com/google/flax/blob/615f40be774e7ed66fd344e8291ac0d48ebcef7d/flax/core/scope.py#L441-L447

The documentation for `rngs` in `apply()` should have a (pointer to) list of names of possible rngs
And documentation for `Dropout` should mention how to pass in rng using `apply()`, without directly passing in like `Dropout()(x,rng=rng)`.
Also probably need to mention the `make_rng()` `fold_in` the rng so each dropout layer will use different rng if there are multiple dropout layers.


Discussion:
We could mention that `Dropout()` requires an `rng` with the name `dropout` in its module documentation. The code is currently very short and it's easily visible, but I agree it would be better discoverable if it was mentioned in the class pydoc as well.

I also think that extending the `Module.apply()` could be extended with something like

```
    rngs: The rngs for the variable collections. For example :class:`flax.linen.stochastic.Dropout`
      requires an additional rng named `dropout`.
```

would be reasonable, since this is quite common and illustrative for newcomers.

WDYT @marcvanzee who was been thinking about the right level of verbosity in our docs

@cccntu why don't you give it a try updating the documentation and sending a PR?
@andsteing  Thanks for the suggestion. I will give it a try in a few days. :)

Resolving patch:
diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
--- a/flax/linen/stochastic.py
+++ b/flax/linen/stochastic.py
@@ -27,6 +27,11 @@
 class Dropout(Module):
   """Create a dropout layer.
 
+    Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
+    to include an RNG seed named `'dropout'`. For example::
+    
+      model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
+
     Attributes:
       rate: the dropout probability.  (_not_ the keep rate!)
       broadcast_dims: dimensions that will share the same dropout mask


Test patch:
Repository: google/flax

Issue:
[Accessibility] Enable EPUB output on ReadTheDocs
Currently there is only HTML output [enabled](https://readthedocs.org/projects/flax/downloads/). It would be great if EPUB and PDF could also be enabled. 


Discussion:
Mind if I do it? It's a really small fix, just gotta add 
```

  - epub
  - pdf
 
```
To the .readthedocs.yml under `formats:`

Resolving patch:
diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
--- a/flax/linen/stochastic.py
+++ b/flax/linen/stochastic.py
@@ -29,7 +29,7 @@ class Dropout(Module):
 
     Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
     to include an RNG seed named `'dropout'`. For example::
-    
+
       model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
 
     Attributes:
diff --git a/flax/struct.py b/flax/struct.py
--- a/flax/struct.py
+++ b/flax/struct.py
@@ -98,7 +98,7 @@ def create(cls, kernel):
   # check if already a flax dataclass
   if '_flax_dataclass' in clz.__dict__:
     return clz
-      
+
   data_clz = dataclasses.dataclass(frozen=True)(clz)
   meta_fields = []
   data_fields = []


Test patch:
diff --git a/tests/struct_test.py b/tests/struct_test.py
--- a/tests/struct_test.py
+++ b/tests/struct_test.py
@@ -67,7 +67,7 @@ def test_keypath_error(self):
       raise e('in_axes')
 
   def test_double_wrap_no_op(self):
-    
+
     class A:
       a: int
Repository: google/flax

Issue:
Make RNG name configurable in Dropout
Currently `'dropout'` is hard-coded as an argument of `make_rng` inside `Dropout`. However, when implementing support for "recurrent dropout"  in an LSTMCell or similar you need two kinds of dropout:
1. A regular dropout which is applied to the inputs with a different mask at each step.
2. A "recurrent dropout" that is applied to the state with the same mask at each step.

To implement `2` a possibility is to set the RNG name to `'recurrent_dropout'` on the Dropout layer applied to the state and guarantee that each step uses the same random state. From `nn.scan`'s perspective the would look like this:

```python
nn.scan(..., split_rngs={'dropout': True, 'recurrent_dropout': False})
```
The proposal is to add an `rng_name` (or similar) attribute to `Dropout` so we are able support these kind of use-cases. The alternative would be to create a separate `RecurrentDropout` layer with the same code but different hard-coded value.


Discussion:
It might be more efficient to compute the mask outside of the loop if you are going to share it. random.* operations are usually quite expensive compared to other operations
Its a good point. Maybe we need a `RNNCellBase.get_stochastic_mask` API and have cells optionally accept a `stochastic_mask` argument.

Resolving patch:
diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
--- a/flax/linen/stochastic.py
+++ b/flax/linen/stochastic.py
@@ -38,10 +38,12 @@ class Dropout(Module):
       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
         masked, whereas if true, no mask is applied and the inputs are returned
         as is.
+      rng_collection: the rng collection name to use when requesting an rng key.
   """
   rate: float
   broadcast_dims: Sequence[int] = ()
   deterministic: Optional[bool] = None
+  rng_collection: str = 'dropout'
 
   @compact
   def __call__(self, inputs, deterministic: Optional[bool] = None):
@@ -67,7 +69,7 @@ def __call__(self, inputs, deterministic: Optional[bool] = None):
     if deterministic:
       return inputs
     else:
-      rng = self.make_rng('dropout')
+      rng = self.make_rng(self.rng_collection)
       broadcast_shape = list(inputs.shape)
       for dim in self.broadcast_dims:
         broadcast_shape[dim] = 1


Test patch:
Repository: google/flax

Issue:
DynamicScale goes to infinity
### System information
- OS Platform and Distribution: Linux 5.15.74-3-MANJARO
- Flax, jax, jaxlib versions: flax 0.6.0, jax 0.3.21, jaxlib 0.3.20+cuda11.cudnn82
- Python version: 3.10.4
- GPU/TPU model and memory: NVIDIA RTX 2060 Super 8GB
- CUDA version (if applicable): cuda11, cudnn82


### Problem you have encountered:
I adapted the gradient scaling code of the imagenet example in order to use fp16 training for my model. During training, the `scale` parameter of `DynamicScaling` became `inf`, because it kept doubling every 2000 steps without the gradients ever becoming `nan`. Once the `scale` hit `inf` it never went down, since dividing `inf` by any number still returns `inf`.

### What you expected to happen:
Either have a `maximum_scale` parameter or add a check to see if `scale` is finite.
I think it's a one line change in the `DynamicScale` code, so I can make a pull request if you agree.

### Steps to reproduce:
```py
import jax.numpy as jnp
from flax.training import dynamic_scale

ds = dynamic_scale.DynamicScale(growth_factor=10.0, growth_interval=1)
x = jnp.float32(1)

for _ in range(70):
    ds, is_fin, aux, grad = ds.value_and_grad(lambda p: p**2)(x)
    x = x - 0.1 * grad
    print(is_fin, ds.scale, x, grad)
```

### Code:
Could check here ([dynamic scale code](https://github.com/google/flax/blob/6b80cbb239b07e370fd0af6655b39ef40de061c0/flax/training/dynamic_scale.py#L139)) if `fin_scale` is finite.


Discussion:
I think we should add an upper bound indeed, are you sure you are really computing in fp16? I'm a bit surprised that this is something you can actually hit in practice
The model parameters are float16, the model inputs are int16 and I accumulate the loss in float32. Now that I think about it, since the loss is accumulated as float32, the grads are float32 (which probably explains why they don't become NaNs). Does it mean that all the backward pass is done in float32?
After checking more extensively, I figured out that part of my model inputs were actually float32. After fixing it the training is about twice as fast, so I wasn't computing in fp16 after all! Thank you for pointing it out.

In any case, it still might be useful to make sure that the scale never goes to infinity.

Resolving patch:
diff --git a/flax/training/dynamic_scale.py b/flax/training/dynamic_scale.py
--- a/flax/training/dynamic_scale.py
+++ b/flax/training/dynamic_scale.py
@@ -132,9 +132,10 @@ def grad_fn_wrapper(*args):
         finite &= jnp.all(lax.is_finite(g))
 
       grow = self.fin_steps == self.growth_interval
-      fin_scale = jnp.where(grow & finite,
-                            self.scale * self.growth_factor,
-                            self.scale)
+      fin_scale = jnp.where(
+          grow & finite,
+          jnp.minimum(self.scale * self.growth_factor, jnp.finfo(jnp.float32).max),
+          self.scale)
       inf_scale = self.scale * self.backoff_factor
       new_scale = jnp.where(finite, fin_scale, inf_scale)
       new_fin_steps = jnp.where(grow | (~finite), 0, self.fin_steps + 1)


Test patch: