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

Issue:
Support serialization for `jax.tree_util.Partial`
`jax.tree_util.Partial` is a PyTree-aware version of `functors.partial` whose leaves (args and keyword args) are arrays.

Therefore, I believe it makes sense to make it `serialization`-aware. This allows to use `Partial` inside of structures that we use for checkpointing.
The definition would be as simple as the following:

```python
# jax.tree_util.Partial does not support flax serialization
# should be upstreamed to Flax
serialization.register_serialization_state(
    jax.tree_util.Partial,
    lambda x: (
        {
            "args": serialization.to_state_dict(x.args),
            "keywords": serialization.to_state_dict(x.keywords),
        }
    ),
    lambda x, sd: jax.tree_util.Partial(
        x.func,
        *serialization.from_state_dict(x.args, sd["args"]),
        **serialization.from_state_dict(x.keywords, sd["keywords"]),
    ),
)
```

Would you accept this to be contributed?


Discussion:
Looks like a nice improvement to me. Feel free to make a PR.

Resolving patch:
diff --git a/flax/serialization.py b/flax/serialization.py
--- a/flax/serialization.py
+++ b/flax/serialization.py
@@ -193,6 +193,20 @@ def _restore_namedtuple(xs, state_dict: Dict[str, Any]):
                              _namedtuple_state_dict,
                              _restore_namedtuple)
 
+register_serialization_state(
+    jax.tree_util.Partial,
+    lambda x: (
+        {
+            "args": to_state_dict(x.args),
+            "keywords": to_state_dict(x.keywords),
+        }
+    ),
+    lambda x, sd: jax.tree_util.Partial(
+        x.func,
+        *from_state_dict(x.args, sd["args"]),
+        **from_state_dict(x.keywords, sd["keywords"]),
+    ),
+)
 
 # On-the-wire / disk serialization format
 


Test patch:
diff --git a/tests/serialization_test.py b/tests/serialization_test.py
--- a/tests/serialization_test.py
+++ b/tests/serialization_test.py
@@ -28,6 +28,7 @@
 from flax.training import train_state
 import jax
 from jax import random
+from jax.tree_util import Partial
 import jax.numpy as jnp
 import msgpack
 import numpy as np
@@ -107,6 +108,16 @@ def test_model_serialization(self):
     restored_model = serialization.from_state_dict(initial_params, state)
     self.assertEqual(restored_model, freeze(state))
 
+  def test_partial_serialization(self):
+    add_one = Partial(jnp.add, 1)
+    state = serialization.to_state_dict(add_one)
+    self.assertEqual(state, {
+        'args': {'0': 1},
+        'keywords': {}
+    })
+    restored_add_one = serialization.from_state_dict(add_one, state)
+    self.assertEqual(add_one.args, restored_add_one.args)
+
   def test_optimizer_serialization(self):
     rng = random.PRNGKey(0)
     module = nn.Dense(features=1, kernel_init=nn.initializers.ones)
Repository: google/flax

Issue:
`typing.get_type_hints()` broken for linen modules
Basically the same issue as #1947, but re-surfaced in `flax==0.6.2`.

This breaks:
```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))
```

This should probably be considered a Python typing design issue; the culprit seems to be [this otherwise valid forward reference](https://github.com/google/flax/blob/3b27a692cf65a48022d4b91c7d6846db74384d12/flax/linen/module.py#L531).

This is easy to work around (by importing `from flax.linen import Module`), but it'd be nice if this wasn't needed!


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
@@ -528,9 +528,6 @@ def reimport(self, other: '_ModuleInternalState') -> None:
 capture_call_intermediates = lambda _, method_name: method_name == '__call__'
 
 
-_ParentType = Union[Type['Module'], Type[Scope], Type[_Sentinel], None]
-
-
 class ParentDescriptor:
   """Wraps parent module references in weak refs.
 
@@ -1651,6 +1648,8 @@ def __call__(self, x):
     return tabulate_fn(*args, **kwargs)
 
 
+_ParentType = Union[Type[Module], Type[Scope], Type[_Sentinel], None]
+
 def merge_param(name: str, a: Optional[T], b: Optional[T]) -> T:
   """Merges construction and call time argument.
 


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
@@ -21,7 +21,7 @@
 import inspect
 import operator
 from typing import (Any, Callable, Generic, Mapping, NamedTuple, Sequence,
-                    Tuple, TypeVar)
+                    Tuple, TypeVar, get_type_hints)
 
 from absl.testing import absltest
 from flax import errors
@@ -1801,6 +1801,13 @@ def test_deepcopy_unspecified_parent(self):
     self.assertIs(unspecified_parent,
                   copy.deepcopy(unspecified_parent))
 
+  def test_type_hints(self):
+    class Network(nn.Module):
+        layers: int
+
+    type_hints = get_type_hints(Network)
+    self.assertEqual(type_hints['layers'], int)
+
 
 class LeakTests(absltest.TestCase):
Repository: google/flax

Issue:
Add two aliases to activations import
Add the JAX standardize alias for normalize as well, see #2088.

This is currently blocked since JAX just added these so we should add them once JAX rolls out a new verison.


Discussion:

Resolving patch:
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -41,6 +41,7 @@
   soft_sign as soft_sign,
   softmax as softmax,
   softplus as softplus,
+  standardize as standardize,
   swish as swish,
   tanh as tanh
 )
diff --git a/flax/linen/activation.py b/flax/linen/activation.py
--- a/flax/linen/activation.py
+++ b/flax/linen/activation.py
@@ -44,6 +44,7 @@
 from jax.nn import soft_sign
 from jax.nn import softmax
 from jax.nn import softplus
+from jax.nn import standardize
 from jax.nn import swish
 import jax.numpy as jnp
 from jax.numpy import tanh


Test patch:
Repository: google/flax

Issue:
Improve documentation for nn.map_variables
Specifically, try to document how to solve the case presented in #2773 by setting `init=self.is_initializing()`. 


Discussion:

Resolving patch:
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -841,21 +841,40 @@ def map_variables(
     methods=None) -> Target:
   """Map Variables inside a module.
 
-  Example::
+  ``map_variables`` can be used to transform the variables inside a module
+  both before and after the module is applied. This is useful among other
+  things for masking the weights of a module without having to modify the
+  module itself.
 
-    class OneBitDense(nn.Module):
-      @nn.compact
-      def __call__(self, x):
-        def sign(x):
-          return jax.tree_util.tree_map(jnp.sign, x)
-        MapDense = nn.map_variables(nn.Dense, "params", sign, init=True)
-        return MapDense(4)(x)
+  Example::
+    >>> import jax
+    >>> import jax.numpy as jnp
+    >>> import flax.linen as nn
+    ...
+    >>> class CasualDense(nn.Module):
+    ...   features: int
+    ...
+    ...   @nn.compact
+    ...   def __call__(self, x):
+    ...     def apply_mask(variables):
+    ...       return (jax.tree_map(jnp.triu, variables)
+    ...               if not self.is_initializing() else variables)
+    ...
+    ...     CasualDense = nn.map_variables(
+    ...       nn.Dense, 'params', apply_mask, init=self.is_initializing())
+    ...
+    ...     mapped_dense = CasualDense(features=self.features, use_bias=False)
+    ...     return mapped_dense(x)
+    ...
+    >>> module = CasualDense(features=5)
+    >>> variables = module.init(jax.random.PRNGKey(0), jnp.ones((1, 5)))
 
   Args:
-    target: the function to be transformed.
+    target: the module or function to be transformed.
     mapped_collections: the collection(s) to be transformed.
-    trans_in_fn: creates a view of the target variables.
-    trans_out_fn: transforms the updated variables in the view after mutation.
+    trans_in_fn: modifies the variables before applying the module or function.
+    trans_out_fn: modifies the variables after applying the module or function,
+      it is only applied if either ``init`` or ``mutable`` are not False.
     init: If True, variables are initialized before transformation.
     mutable: If True, the mapped variable collections will be mutable.
     rngs: PRNGSequences added to the transformed scope (default: all).


Test patch:
Repository: google/flax

Issue:
Cannot import flax.training.checkpoints in 0.6.4
### System information
- OS Platform and Distribution: Ubuntu 22.04.1 LTS, also in Colab environment
- Flax, jax, jaxlib versions:
  * flax 0.6.4
  * jax 0.3.25
  * jaxlib 0.3.25
- Python version: 3.10.6
- GPU/TPU model and memory: No Accelerator / 16GB

### Problem you have encountered:
With FLAX v0.6.4 I can't import `flax.training.checkpoints` module due to following error:
```
ImportError: cannot import name 'monitoring' from 'jax' (/usr/local/lib/python3.8/dist-packages/jax/__init__.py)
```
This does not happen in v0.6.3.

### What you expected to happen:
The module should be imported.

### Logs, error messages, etc:
Error message from jupyter notebook:
```
ImportError                               Traceback (most recent call last)

[<ipython-input-3-9a234296e658>](https://localhost:8080/#) in <module>
      1 import flax
----> 2 from flax.training import checkpoints

[/usr/local/lib/python3.8/dist-packages/flax/training/checkpoints.py](https://localhost:8080/#) in <module>
     36 from flax import traverse_util
     37 import jax
---> 38 from jax import monitoring
     39 from jax import process_index
     40 from jax import sharding

ImportError: cannot import name 'monitoring' from 'jax' (/usr/local/lib/python3.8/dist-packages/jax/__init__.py)
```

### Steps to reproduce:
[Colab notebook](https://colab.research.google.com/drive/1ZLR1JSJPfaaoTmL7bow8oebqyhhxrqSo?usp=sharing)



Discussion:
The problem seems to be that `jax.monitoring` is available since `0.4` but flax requires `>=0.3.16`. I was also able to resolve this problem by upgrading `jax` and `jaxlib` to `0.4.2`. 
`!python -m pip install jax==0.4.2 jaxlib==0.4.2`

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.16",
+    "jax>=0.4.2",
     "matplotlib",  # only needed for tensorboard export
     "msgpack",
     "optax",


Test patch:
Repository: google/flax

Issue:
Improve documentation of flax.nn.remat for multiple methods
### Problem you have encountered:

I found the ```flax.linen.remat``` [docs](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.remat.html) a little confusing when it comes to checkpointing multiple methods on a module.

### What you expected to happen:

I expected the docs to tell me the expected type for ```flax.linen.remat```'s "methods" argument. 

I initially assumed it worked similarly to ```Module.apply```, which expects the argument ```method``` to be "a function to call apply on". However, it [turned out](https://github.com/google/flax/blob/e51d01734e1264bf54e64ee28ce1bef7fb617131/flax/linen/transforms.py#L298) that ```flax.linen.remat``` expected the argument ```methods``` be a list of strings. 


Discussion:

Resolving patch:
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -588,7 +588,30 @@ def checkpoint(target: Target,
                methods=None) -> Target:
   """Lifted version of ``jax.checkpoint``.
 
-  This function is aliased to ``lift.remat`` just like ``jax.remat``.
+  Checkpointing is a technique for reducing memory usage by recomputing
+  activations during backpropagation. When training large models, it can be
+  helpful to checkpoint parts of the model to trade off memory usage for
+  additional computation.
+
+  Example::
+
+    >>> import jax
+    >>> import jax.numpy as jnp
+    >>> import flax.linen as nn
+    ...
+    >>> class CheckpointedMLP(nn.Module):
+    ...   @nn.compact
+    ...   def __call__(self, x):
+    ...     CheckpointDense = nn.checkpoint(nn.Dense)
+    ...     x = CheckpointDense(128)(x)
+    ...     x = nn.relu(x)
+    ...     x = CheckpointDense(1)(x)
+    ...     return x
+    ...
+    >>> model = CheckpointedMLP()
+    >>> variables = model.init(jax.random.PRNGKey(0), jnp.ones((1, 16)))
+
+  This function is aliased to ``remat`` just like ``jax.remat``.
 
   Args:
     target: a ``Module`` or a function taking a ``Module``
@@ -616,7 +639,9 @@ def checkpoint(target: Target,
       arguments as static can avoid ConcretizationTypeErrors when tracing, but
       at the cost of more retracing overheads.
     policy: Experimental checkpoint policy, see ``jax.checkpoint``.
-    methods: If `target` is a `Module`, the methods of `Module` to checkpoint.
+    methods: An optional list of method names that will be lifted, if `methods`
+      is None (default) only the `__call__` method will be lifted. If `target`
+      is a function, `methods` is ignored.
 
   Returns:
     A wrapped version of ``target``. When computing gradients intermediate


Test patch:
Repository: google/flax

Issue:
Show the groups for the flax.linen API references TOC



Discussion:

Resolving patch:
diff --git a/docs/_ext/flax_module.py b/docs/_ext/flax_module.py
new file mode 100644
--- /dev/null
+++ b/docs/_ext/flax_module.py
@@ -0,0 +1,76 @@
+# Copyright 2023 The Flax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sphinx directive for visualizing Flax modules.
+
+Use directive as follows:
+
+.. flax_module::
+  :module: flax.linen
+  :class: Dense
+
+"""
+
+from docutils import nodes
+from docutils.parsers.rst import directives
+from docutils.statemachine import ViewList
+
+import sphinx
+from sphinx.util.docutils import SphinxDirective
+from docs.conf_sphinx_patch import generate_autosummary_content
+import sphinx.ext.autosummary.generate as ag
+import importlib
+
+
+def render_module(modname: str, qualname: str, app):
+  parent = importlib.import_module(modname)
+  obj = getattr(parent, qualname)
+  template = ag.AutosummaryRenderer(app)
+  template_name = "flax_module"
+  imported_members = False
+  recursive = False
+  context = {}
+  return generate_autosummary_content(
+    qualname, obj, parent, template, template_name, imported_members,
+    app, recursive, context, modname, qualname)
+
+class FlaxModuleDirective(SphinxDirective):
+  has_content = True
+  option_spec = {
+    'module': directives.unchanged,
+    'class': directives.unchanged,
+  }
+
+  def run(self):
+    module_template = render_module(
+      self.options['module'], self.options['class'], self.env.app
+    )
+    module_template = module_template.splitlines()
+
+    # Create a container for the rendered nodes
+    container_node = nodes.container()
+    self.content = ViewList(module_template, self.content.parent)
+    self.state.nested_parse(self.content, self.content_offset, container_node)
+
+    return [container_node]
+
+
+def setup(app):
+  app.add_directive('flax_module', FlaxModuleDirective)
+
+  return {
+      'version': sphinx.__display_version__,
+      'parallel_read_safe': True,
+      'parallel_write_safe': True,
+  }
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -60,6 +60,7 @@
     'sphinx.ext.viewcode',
     'myst_nb',
     'codediff',
+    'flax_module',
     'sphinx_design',
 ]
 


Test patch:
Repository: google/flax

Issue:
ResNetV1 implementation is actually ResNetV1.5
[Downsampling is performed by the 3x3 convolution](https://github.com/google/flax/blob/3ea6381139a340fa7b29c74a84205d8bad18b489/examples/imagenet/models.py#L74), which means it is actually [ResNetV1.5](https://catalog.ngc.nvidia.com/orgs/nvidia/resources/resnet_50_v1_5_for_pytorch), not ResNetV1 as claimed. (In the original version the first convolution in the block has stride 2: see [original paper](https://arxiv.org/pdf/1512.03385.pdf), the caption to Table 1.)


Discussion:

Resolving patch:
diff --git a/examples/imagenet/models.py b/examples/imagenet/models.py
--- a/examples/imagenet/models.py
+++ b/examples/imagenet/models.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Flax implementation of ResNet V1."""
+"""Flax implementation of ResNet V1.5."""
 
 # See issue #620.
 # pytype: disable=wrong-arg-count
@@ -87,7 +87,7 @@ def __call__(self, x):
 
 
 class ResNet(nn.Module):
-  """ResNetV1."""
+  """ResNetV1.5."""
 
   stage_sizes: Sequence[int]
   block_cls: ModuleDef


Test patch:
diff --git a/examples/imagenet/models_test.py b/examples/imagenet/models_test.py
--- a/examples/imagenet/models_test.py
+++ b/examples/imagenet/models_test.py
@@ -26,11 +26,11 @@
 jax.config.update('jax_disable_most_optimizations', True)
 
 
-class ResNetV1Test(parameterized.TestCase):
-  """Test cases for ResNet v1 model definition."""
+class ResNetTest(parameterized.TestCase):
+  """Test cases for ResNet v1.5 model definition."""
 
-  def test_resnet_v1_model(self):
-    """Tests ResNet V1 model definition and output (variables)."""
+  def test_resnet_model(self):
+    """Tests ResNet V1.5 model definition and output (variables)."""
     rng = jax.random.key(0)
     model_def = models.ResNet50(num_classes=10, dtype=jnp.float32)
     variables = model_def.init(rng, jnp.ones((8, 224, 224, 3), jnp.float32))
@@ -43,8 +43,8 @@ def test_resnet_v1_model(self):
     self.assertLen(variables['params'], 19)
 
   @parameterized.product(model=(models.ResNet18, models.ResNet18Local))
-  def test_resnet_18_v1_model(self, model):
-    """Tests ResNet18 V1 model definition and output (variables)."""
+  def test_resnet_18_model(self, model):
+    """Tests ResNet18 V1.5 model definition and output (variables)."""
     rng = jax.random.key(0)
     model_def = model(num_classes=2, dtype=jnp.float32)
     variables = model_def.init(rng, jnp.ones((1, 64, 64, 3), jnp.float32))
Repository: google/flax

Issue:
`Module.tabulate` does not support `Enum` inputs
An internal request to add `Enum` input support for `Module.tabulate`
```
class Net(nn.Module):
  @nn.compact
  def __call__(self, inputs):
    x = inputs['x']
    x = nn.Dense(features=2)(x)
    return jnp.sum(x)

inputs = {'x': jnp.ones((1, 1))}
print(Net().tabulate(jax.random.key(0), inputs=inputs)) # this input works fine

class InputEnum(str, enum.Enum):
  x = 'x'

inputs = {InputEnum.x: jnp.ones((1, 1))}
print(Net().tabulate(jax.random.key(0), inputs)) # this Enum input throws an error
```
```
RepresenterError: ('cannot represent an object', <InputEnum.x: 'x'>)
```


Discussion:

Resolving patch:
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -15,6 +15,7 @@
 """Flax Module summary library."""
 
 import dataclasses
+import enum
 import io
 from abc import ABC, abstractmethod
 from types import MappingProxyType
@@ -709,12 +710,17 @@ def _normalize_structure(obj):
   if isinstance(obj, (tuple, list)):
     return tuple(map(_normalize_structure, obj))
   elif isinstance(obj, Mapping):
-    return {k: _normalize_structure(v) for k, v in obj.items()}
+    return {
+      _normalize_structure(k): _normalize_structure(v) for k, v in obj.items()
+    }
   elif dataclasses.is_dataclass(obj):
     return {
       f.name: _normalize_structure(getattr(obj, f.name))
       for f in dataclasses.fields(obj)
     }
+  elif isinstance(obj, enum.Enum):
+    # `yaml.safe_dump` does not support Enum key types so extract the underlying value
+    return obj.value
   else:
     return obj
 


Test patch:
diff --git a/tests/linen/summary_test.py b/tests/linen/summary_test.py
--- a/tests/linen/summary_test.py
+++ b/tests/linen/summary_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
 from typing import List
 
 import jax
@@ -735,6 +736,26 @@ def __call__(self, x):
     lines = rep.splitlines()
     self.assertIn('Total Parameters: 50', lines[-2])
 
+  def test_tabulate_enum(self):
+    class Net(nn.Module):
+      @nn.compact
+      def __call__(self, inputs):
+        x = inputs['x']
+        x = nn.Dense(features=2)(x)
+        return jnp.sum(x)
+
+    class InputEnum(str, enum.Enum):
+      x = 'x'
+
+    inputs = {InputEnum.x: jnp.ones((1, 1))}
+    # test args
+    lines = Net().tabulate(jax.random.key(0), inputs).split('\n')
+    self.assertIn('x: \x1b[2mfloat32\x1b[0m[1,1]', lines[5])
+    # test kwargs
+    lines = Net().tabulate(jax.random.key(0), inputs=inputs).split('\n')
+    self.assertIn('inputs:', lines[5])
+    self.assertIn('x: \x1b[2mfloat32\x1b[0m[1,1]', lines[6])
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/flax

Issue:
Improve SEO for docs pages
When I Google a Flax module to pull up its API documentation, I get broken or irrelevant pages. Examples:

- [flax dense](https://www.google.com/search?q=flax%20dense) | [top result](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.Dense.html) | [desired page](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.Dense)
- [flax layernorm](https://www.google.com/search?q=flax%20layernorm) | [top result](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.LayerNorm.html) | [desired page](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.LayerNorm)
- [flax multi head attention](https://www.google.com/search?q=flax%20attention) | [top result](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.MultiHeadDotProductAttention.html) | [desired page](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.MultiHeadAttention)
- [flax conv](https://www.google.com/search?q=flax%20conv) | [top result](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.Conv.html) | [desired page](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.Conv)

Is there something that can be done on the website side to improve this?


Discussion:

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -108,6 +108,8 @@
 # so a file named 'default.css' will overwrite the builtin 'default.css'.
 html_static_path = ['_static']
 
+html_extra_path = ['robots.txt']
+
 html_theme_options = {
   'repository_url': 'https://github.com/google/flax',
   'use_repository_button': True,  # add a 'link to repository' button


Test patch: