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

Issue:
Pooling: passing "sequence of `n` `(low, high)` integer pairs" resulting in TypeError
Trying to pass a tuple or list of tuples to a pool operation's padding parameter gives out the following errors: 
`TypeError: Unknown padding type: (1, 1).`
`TypeError : unhashable type: 'list' `


Sample code for reproducing the bug:
```python3
from flax import nn
from jax import random

class FlaxModel(nn.Module):
    def apply(self, x):
        x = nn.max_pool(x, (3, 3), strides=(2, 2), padding=[(1, 1), (1, 1)])
        return x

rng = random.PRNGKey(0)
model, _ = FlaxModel.init_by_shape(rng, [(1, 100, 100, 1)])
```



Discussion:
Indeed looks like our code doesn't support padding that's a sequence of pairs. @hawkinsp has said that the version of JAX on HEAD added support for this, we should add a test and plumb it through correctly. (Or in the meanwhile if that's impossible, support this by manually padding before calling into `lax.reduce_window`)
Yes, JAX at head supports a sequence of `(low, high)` padding pairs.

Flax probably still needs to do some work to add batch and feature dimensions to what the user provides. JAX and XLA don't have opinions about which dimensions are batch and which are feature, but Flax is documented to only accept padding for the spatial dimensions.

Resolving patch:
diff --git a/flax/nn/pooling.py b/flax/nn/pooling.py
--- a/flax/nn/pooling.py
+++ b/flax/nn/pooling.py
@@ -44,6 +44,14 @@ def pool(inputs, init, reduce_fn, window_shape, strides, padding):
   strides = strides or (1,) * len(window_shape)
   strides = (1,) + strides + (1,)
   dims = (1,) + window_shape + (1,)
+  if not isinstance(padding, str):
+    padding = tuple(map(tuple, padding))
+    assert(len(padding) == len(window_shape)), (
+      f"padding {padding} must specify pads for same number of dims as "
+      f"window_shape {window_shape}")
+    assert(all([len(x) == 2 for x in padding])), (
+      f"each entry in padding {padding} must be length 2")
+    padding = ((0,0),) + padding + ((0,0),)
   return lax.reduce_window(inputs, init, reduce_fn, dims, strides, padding)
 
 


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -545,6 +545,24 @@ def test_max_pool(self):
     ]).reshape((1, 3, 3, 1))
     onp.testing.assert_allclose(y_grad, expected_grad)
 
+  def test_max_pool_explicit_pads(self):
+    x = jnp.arange(9).reshape((1, 3, 3, 1)).astype(jnp.float32)
+    pool = lambda x: nn.max_pool(x, (2, 2), padding=((1,1),(1,1)))
+    expected_y = jnp.array([
+        [0.,1.,2.,2.],
+        [3.,4.,5.,5.],
+        [6.,7.,8.,8.],
+        [6.,7.,8.,8.],
+    ]).reshape((1, 4, 4, 1))
+    y = pool(x)
+    onp.testing.assert_allclose(y, expected_y)
+    y_grad = jax.grad(lambda x: pool(x).sum())(x)
+    expected_grad = jnp.array([
+        [1., 1., 2.],
+        [1., 1., 2.],
+        [2., 2., 4.],
+    ]).reshape((1, 3, 3, 1))
+    onp.testing.assert_allclose(y_grad, expected_grad)
 
 class NormalizationTest(absltest.TestCase):
Repository: google/flax

Issue:
Syntax warning due to comparison of literals using is in Python 3.8
### Problem you have encountered:

Syntax warning due to comparison of literals using is in Python 3.8 to use != .

### Steps to reproduce:

```
find . -iname '*.py' | grep -v example | grep -v doc | xargs -P4 -I{} python3.8 -Wall -m py_compile {}
./flax/optim/weight_norm.py:150: SyntaxWarning: "is not" with a literal. Did you mean "!="?
  if decay is not 0:
```


Discussion:

Resolving patch:
diff --git a/flax/optim/weight_norm.py b/flax/optim/weight_norm.py
--- a/flax/optim/weight_norm.py
+++ b/flax/optim/weight_norm.py
@@ -147,7 +147,7 @@ def _split_grad(self, param, state, grad, decay):
       scale_grad = jnp.sum(
           grad * direction, axis=red_dims, keepdims=True)
       direction_grad = state.mult * (grad - scale_grad * direction)
-      if decay is not 0:
+      if decay != 0:
         direction_grad = direction_grad + decay * direction
       direction_info = direction, state.direction_state, direction_grad
       scale_info = scale, state.scale_state, scale_grad


Test patch:
Repository: google/flax

Issue:
PPO example does not terminate properly
### Configuration

Running the PPO example for a short number of frames  in order to reproduce as fast as possible on a cloud VM with a V100 GPU. Config python3.7, flax 0.2.2, jax 0.2.1, jaxlib 0.1.55 .

Command run:
`python ppo_main.py --config.game=Qbert --config.total_frames=4000`

### Problem you have encountered:

Program does not exit. One can `print('Done')` after `ppo_lib.train` in `ppo_main` but there is an open thread and program can't exit (even after adding `raise SystemExit`).

### Extra comments

Added extra line in `main`  ` tf.config.experimental.set_visible_devices([],'GPU')` in order for the program to run properly with `tensorflow-gpu`, this is common in other `flax/examples`. 


Discussion:

Resolving patch:
diff --git a/examples/ppo/agent.py b/examples/ppo/agent.py
--- a/examples/ppo/agent.py
+++ b/examples/ppo/agent.py
@@ -43,6 +43,7 @@ def __init__(self, game: str):
     parent_conn, child_conn = multiprocessing.Pipe()
     self.proc = multiprocessing.Process(
         target=rcv_action_send_exp, args=(child_conn, game))
+    self.proc.daemon = True
     self.conn = parent_conn
     self.proc.start()
 
diff --git a/examples/ppo/ppo_main.py b/examples/ppo/ppo_main.py
--- a/examples/ppo/ppo_main.py
+++ b/examples/ppo/ppo_main.py
@@ -19,6 +19,8 @@
 import jax.random
 from ml_collections import config_flags
 
+import tensorflow as tf
+
 import ppo_lib
 import models
 import env_utils
@@ -34,6 +36,9 @@
     'File path to the default configuration file.')
 
 def main(argv):
+  # Make sure tf does not allocate gpu memory.
+  tf.config.experimental.set_visible_devices([], 'GPU')
+
   config = FLAGS.config
   game = config.game + 'NoFrameskip-v4'
   num_actions = env_utils.get_num_actions(game)


Test patch:
Repository: google/flax

Issue:
QoL: better print for FrozenDict
The best way I'm aware of to get an overview of model shape is via `jax.tree_map(jnp.shape, params)`. FrozenDicts have no concept of pretty printing the way dicts do, so large models are unwieldy to parse at a glance. 


Discussion:
Yes I noticed the output ends up without indentst and newlines. Let's try to fix that.

Resolving patch:
diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py
--- a/flax/core/frozen_dict.py
+++ b/flax/core/frozen_dict.py
@@ -24,6 +24,14 @@
 V = TypeVar('V')
 
 
+def _indent(x, num_spaces):
+  indent_str = ' ' * num_spaces
+  lines = x.split('\n')
+  assert lines[-1] == ''
+  # skip the final line because it's empty and should not be indented.
+  return '\n'.join(indent_str + line for line in lines[:-1]) + '\n'
+
+
 @jax.tree_util.register_pytree_node_class
 class FrozenDict(Mapping[K, V]):
   """An immutable variant of the Python dict."""
@@ -55,7 +63,21 @@ def __len__(self):
     return len(self._dict)
 
   def __repr__(self):
-    return 'FrozenDict(%r)' % self._dict
+    return self.pretty_repr()
+
+  def pretty_repr(self, num_spaces=4):
+    """Returns an indented representation of the nested dictionary."""
+    def pretty_dict(x):
+      if not isinstance(x, dict):
+        return repr(x)
+      rep = ''
+      for key, val in x.items():
+        rep += f'{key}: {pretty_dict(val)},\n'
+      if rep:
+        return '{\n' + _indent(rep, num_spaces) + '}'
+      else:
+        return '{}'
+    return f'FrozenDict({pretty_dict(self._dict)})'
 
   def __hash__(self):
     if self._hash is None:


Test patch:
diff --git a/tests/core/frozen_dict_test.py b/tests/core/frozen_dict_test.py
--- a/tests/core/frozen_dict_test.py
+++ b/tests/core/frozen_dict_test.py
@@ -59,5 +59,20 @@ def test_frozen_items(self):
     self.assertEqual(items, [('a', 1), ('b', freeze(xs['b']))])
 
 
+  def test_frozen_dict_repr(self):
+    expected = (
+"""FrozenDict({
+    a: 1,
+    b: {
+        c: 2,
+        d: {},
+    },
+})""")
+
+    xs = FrozenDict({'a': 1, 'b': {'c': 2, 'd': {}}})
+    self.assertEqual(repr(xs), expected)
+    self.assertEqual(repr(FrozenDict()), 'FrozenDict({})')
+
+
 if __name__ == '__main__':
   absltest.main()
Repository: google/flax

Issue:
After update from 0.2.0: AttributeError: module 'jax.core' has no attribute 'eval_context'
After updating from flax 0.2.0 to flax 0.2.2 I get the above error message. Downgrading to 0.2.0 solves this, so the error source is located. I'm working with the now deprecated flax.nn package if backward-compatibility might be the reason for this issue.
The Issue is encountered in a custom RNN, when using the init_by_shape function in conjunction with jax.lax.scan.


Discussion:
Hi @mr128254 -- we are about to officially deprecate `flax.nn` but if you have a minimal repro we can perhaps take a look. (Also have you upgraded your version of JAX in parallel to upgrading the Flax version?)
I am pretty sure this has something to do with the Jax version. We should raise the minimal version in `setup.py`

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.1.59",
+    "jax>=0.1.77",
     "matplotlib",  # only needed for tensorboard export
     "dataclasses;python_version<'3.7'", # will only install on py3.6
     "msgpack",


Test patch:
Repository: google/flax

Issue:
Improve Documentation: Jitting init()
In some of our examples we `jax.jit` the `init()`, for instance in the [WMT example](https://github.com/google/flax/blob/master/linen_examples/wmt/train.py#L472), and in other examples we don't.

@Marvin182 mentioned in a chat: "Jitting the model.init() of the PixelCNN example takes forever (>5min) but runs without jitting in <20s. " [on TPU]

@jheek replied: "jitting init is a difficult thing. On one hand we save time because we are lazy and potentially avoid lot's of small compiles. on the other hand we have a lot of duplicate HLO's when a model has many paramaters with the same shape & dtype"


It thus seems there are some best practices on how and when to `jit` the `init()` of Flax modules, and it would be useful to document this since it can make a big difference in practice, especially on TPU.


Discussion:
Another comment from Daniel Johnson:

"if you expect to create some expensive value and then immediately throw it away. In the case of flax models, if you call model.init it actually runs a forward pass through the model, which can be expensive. Putting it in jit means jax will trace through it, then dead-code-eliminate the forward pass and only keep the parameters."

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -725,6 +725,13 @@ def init(self, rngs: Union[PRNGKey, RNGSequences], *args,
            **kwargs) -> VariableDict:
     """Initializes a module method with variables and returns modified variables.
 
+    Jitting `init` initializes a model lazily using only the shapes of the 
+    provided arguments, and avoids computing the forward pass with actual 
+    values. Example::
+
+      jit_init = jax.jit(SomeModule.init)
+      jit_init(rng, jnp.ones(input_shape, jnp.float32))      
+
     Args:
       rngs: The rngs for the variable collections.
       method: An optional method. If provided, applies this method. If not


Test patch:
Repository: google/flax

Issue:
When jax_enable_x64 is set Adam promotes everything to float64

### Problem you have encountered:

When `jax_enable_x64` is set, Adam's `apply_gradient` method will promote all float32 arrays to float64, potentially unexpectedly degrading performance.

This is due to jax's wonky type promotion semantics. The offending line is:
https://github.com/google/flax/blob/3e36db3e5e3b8e6e1777d612f270e7948238aa9c/flax/optim/adam.py#L82

which promotes like:
```python
jnp.array([0], dtype=jnp.int32) + 1. # == DeviceArray([1.], dtype=float64)
```
and then cascades from there promoting everything to float64

### What you expected to happen:

Arrays should retain their dtypes on optimizer updates.

### Logs, error messages, etc:


### Steps to reproduce:

```python
from jax.config import config
config.update("jax_enable_x64", True)

import jax.numpy as jnp
import flax

opt = flax.optim.Adam(1e-3).create(
        {"x": jnp.zeros(10, dtype=jnp.float32)})

assert opt.target["x"].dtype == jnp.float32

opt = opt.apply_gradient({"x": jnp.zeros(10, dtype=jnp.float32)})

# This fails, since dtype was promoted to float64
assert opt.target["x"].dtype == jnp.float32
```


Discussion:

Resolving patch:
diff --git a/flax/optim/adam.py b/flax/optim/adam.py
--- a/flax/optim/adam.py
+++ b/flax/optim/adam.py
@@ -98,7 +98,7 @@ def apply_param_gradient(self, step, hyper_params, param, state, grad):
     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq
 
     # bias correction
-    t = step + 1.
+    t = jnp.array(step + 1, lax.dtype(param.dtype))
     grad_ema_corr = grad_ema / (1 - beta1 ** t)
     grad_sq_ema_corr = grad_sq_ema / (1 - beta2 ** t)
 
diff --git a/flax/optim/lamb.py b/flax/optim/lamb.py
--- a/flax/optim/lamb.py
+++ b/flax/optim/lamb.py
@@ -74,7 +74,7 @@ def apply_param_gradient(self, step, hyper_params, param, state, grad):
     grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad
     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq
 
-    t = step + 1.
+    t = jnp.array(step + 1, lax.dtype(param.dtype))
     grad_ema_corr = grad_ema / (1. - beta1 ** t)
     grad_sq_ema_corr = grad_sq_ema / (1. - beta2 ** t)
 


Test patch:
Repository: google/flax

Issue:
Port ensembling HOWTO from old diff based system
And instead, use a standalone doc with tests like in #771

Here is the old (pre-Linen) HOWTO diff, for reference:
https://github.com/google/flax/blob/master/howtos/diffs/ensembling.diff


Discussion:

Resolving patch:
diff --git a/docs/_ext/codediff.py b/docs/_ext/codediff.py
--- a/docs/_ext/codediff.py
+++ b/docs/_ext/codediff.py
@@ -26,14 +26,14 @@
 Use directive as follows:
 
 .. codediff::
-  :title-left: <LEFT_CODE_BLOCK_TITLE>
-  :title-right: <RIGHT_CODE_BLOCK_TITLE>
-  :highlight-left: <LINES_TO_HIGHLIGHT_LEFT>
-  :highlight-right: <LINES_TO_HIGHLIGHT_RIGHT>
+  :title_left: <LEFT_CODE_BLOCK_TITLE>
+  :title_right: <RIGHT_CODE_BLOCK_TITLE>
   
   <CODE_BLOCK_LEFT>
   ---
   <CODE_BLOCK_RIGHT>
+
+In order to highlight a line of code, prepend it with "#!".
 """
 
 class CodeDiffParser:
@@ -94,7 +94,7 @@ class CodeDiffDirective(SphinxDirective):
     'code_sep': directives.unchanged,
     }
 
-  def run(self):    
+  def run(self):
     new_content = CodeDiffParser().parse(list(self.content), **self.options)
 
     node = nodes.paragraph()


Test patch:
Repository: google/flax

Issue:
 threading.Condition.notifyAll has been deprecated in favour of notify_all in Python 3.10
### Problem you have encountered:

 `threading.Condition.notifyAll` has been deprecated in favour of `notify_all` in Python 3.10. Ref : python/cpython#25174

### What you expected to happen:

use `notify_all` in below places.

```
rg -t py -w 'currentThread|notifyAll|activeCount|isDaemon|setDaemon' 
flax/training/prefetch_iterator.py
58:        self._cond.notifyAll()
68:      self._cond.notifyAll()
80:          self._cond.notifyAll()
88:          self._cond.notifyAll()
```


Discussion:

Resolving patch:
diff --git a/flax/training/prefetch_iterator.py b/flax/training/prefetch_iterator.py
--- a/flax/training/prefetch_iterator.py
+++ b/flax/training/prefetch_iterator.py
@@ -55,7 +55,7 @@ def __next__(self):
       self._cond.wait_for(lambda: self._buffer or not self._active)
       if self._buffer:
         item = self._buffer.pop(0)
-        self._cond.notifyAll()
+        self._cond.notify_all()
         return item
       if self._error:
         raise self._error  # pylint: disable=raising-bad-type
@@ -65,7 +65,7 @@ def __next__(self):
   def close(self):
     with self._cond:
       self._active = False
-      self._cond.notifyAll()
+      self._cond.notify_all()
 
   def _prefetch_loop(self):
     """Prefetch loop that prefetches a tf dataset."""
@@ -77,7 +77,7 @@ def _predicate():
         item = next(self._data_iter)
         with self._cond:
           self._buffer.append(item)
-          self._cond.notifyAll()
+          self._cond.notify_all()
           self._cond.wait_for(_predicate)
           if not self._active:
             return
@@ -85,5 +85,5 @@ def _predicate():
         with self._cond:
           self._error = e
           self._active = False
-          self._cond.notifyAll()
+          self._cond.notify_all()
           return


Test patch:
Repository: google/flax

Issue:
AttributeError: module 'flax.linen' has no attribute 'merge_param'
[This guide](https://flax.readthedocs.io/en/latest/design_notes/arguments.html) suggests using `nn.merge_param` to combine arguments, but `merge_param` is only available through `nn.module.merge_param`. I believe it needs to be added to the import line [here](https://github.com/google/flax/blob/4ae9143f7ef46ffab6d9123ba1b2e4f3303e68d1/flax/linen/__init__.py#L28). I can open a PR if this is the case.



Discussion:
Good catch! Please do open that PR

Resolving patch:
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -25,7 +25,8 @@
                         make_causal_mask, combine_masks)
 from ..core import broadcast, DenyList
 from .linear import Conv, ConvTranspose, Dense, DenseGeneral, Embed
-from .module import Module, compact, enable_named_call, disable_named_call, Variable, init, init_with_output, apply
+from .module import (Module, compact, enable_named_call, disable_named_call,
+                     Variable, init, init_with_output, apply, merge_param)
 from .normalization import BatchNorm, GroupNorm, LayerNorm
 from .pooling import avg_pool, max_pool
 from .recurrent import GRUCell, LSTMCell, ConvLSTM, OptimizedLSTMCell


Test patch:
Repository: google/flax

Issue:
flax.core.FrozenDict copy broken when the new dictionary contains some names
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### Problem you have encountered:
Adding a dictionary which contains 'cls' key fails, 
![image](https://user-images.githubusercontent.com/6980056/125225456-33f22600-e284-11eb-8b32-8ef1bae5ac3f.png)

### What you expected to happen:
expected to update the value of 'cls' key. 

### Logs, error messages, etc:



### Steps to reproduce:

```
flax.core.FrozenDict({}).copy({'cls': 'abc'})
```

One way to workaround this is to manually create concatenated FrozenDict instead of using `copy`.
```
flax.core.FrozenDict({**flax.core.FrozenDict({'def': '123', 'cls': 22}), **{'cls': 'abc'}})
```


Discussion:
Thanks for catching this bug!

Your code snippet on line 98 resolves to: `return FrozenDict(self, cls='abc')`. This will invoke `__new__` of a superclass in the `Typing` library, which has `cls` as its first argument. So since you also pass it as a kwarg, the interpreter complains that you passed the same argument twice.

It seems dangerous to me that the current code just expands all key-value pairs in `add_or_replace` to kwargs to the constructor, since any reserved word could introduce bugs. The safest way seems to me to explicitly wrap the two dicts in a new dict, i.e. replace line 98 with:

```python
return type(self)({**self, **unfreeze(add_or_replace)})
```

@jheek WDYT?
Oh my Python! :)

@marcvanzee your solution looks like the easiest workaround

Resolving patch:
diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py
--- a/flax/core/frozen_dict.py
+++ b/flax/core/frozen_dict.py
@@ -95,7 +95,7 @@ def __hash__(self):
 
   def copy(self, add_or_replace: Mapping[K, V]) -> 'FrozenDict[K, V]':
     """Create a new FrozenDict with additional or replaced entries."""
-    return type(self)(self, **unfreeze(add_or_replace))
+    return type(self)({**self, **unfreeze(add_or_replace)})
 
   def items(self):
     for key in self._dict:


Test patch:
diff --git a/tests/core/frozen_dict_test.py b/tests/core/frozen_dict_test.py
--- a/tests/core/frozen_dict_test.py
+++ b/tests/core/frozen_dict_test.py
@@ -80,6 +80,10 @@ def test_frozen_dict_reduce(self):
     self.assertEqual(before, after)
     self.assertEqual(after, {'a': {'b': 1, 'c': 2}})
 
+  def test_frozen_dict_copy_reserved_name(self):
+    result = FrozenDict({'a': 1}).copy({'cls': 2})
+    self.assertEqual(result, {'a': 1, 'cls': 2})
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/flax

Issue:
logical_to_mesh_axes does not process rules with repeated array dim names correctly.

### Problem you have encountered:
Current implementation of logical_to_mesh_axes function results in incorrect annotation propagation incase logical_axis_rules have more than one entry for a logical axis. For example:
 logical_axis_rules = (('batch', 'data'),
      ('vocab', 'model'),
      ('mlp', 'model'),
      ('heads', 'model'),
      ('joined_kv', None),
      ('kv', None),
      ('embed', 'model'),
      ('embed', 'data'),
      ('relpos_buckets', None),
      ('length', None),
      ('layers', None),
      ('stack', None),
    )

should annotate the following tensor:

y = with_sharding_constraint(y, ('batch', 'length', 'embed'))

to 
 
axis_resources=<partitions=((\'data\',), (), (\'model\',)) in the resulting pre optimization HLO.

However with the current fuction it results in:

axis_resources=<partitions=((\'data\',), (), ()).

### Steps to reproduce:
logical_axis_rules, sharding constraints mentioned above should suffice to repro the issue.

The issue seems to be at:
https://github.com/google/flax/blob/d10eda85791d5cb9029cf431aae10c7032c2ea8b/flax/linen/partitioning.py#L166


Discussion:

Resolving patch:
diff --git a/flax/linen/partitioning.py b/flax/linen/partitioning.py
--- a/flax/linen/partitioning.py
+++ b/flax/linen/partitioning.py
@@ -164,7 +164,8 @@ def logical_to_mesh_axes(array_dim_names: Sequence[str],
     if rule_model_name in array_dim_names:
       pos = array_dim_names.index(rule_model_name)
       if rule_mesh_name is None or rule_mesh_name in result:
-        result[pos] = None
+        if result[pos] == _unassigned_axis:
+          result[pos] = None
       else:
         result[pos] = result[pos] or rule_mesh_name
   if _unassigned_axis in result:


Test patch:
diff --git a/tests/linen/partitioning_test.py b/tests/linen/partitioning_test.py
--- a/tests/linen/partitioning_test.py
+++ b/tests/linen/partitioning_test.py
@@ -72,6 +72,16 @@ def test_logical_to_mesh_axes(self):
     with partitioning.axis_rules(AXIS_RULES_1):
       with self.assertRaises(ValueError):
         partitioning.logical_to_mesh_axes(('foo', 'foo', 'baz'))
+  def test_logical_to_mesh_axes_overrides(self):
+    p_rules = (
+        ('baz', 'data'),
+        ('bar', None),
+        ('foo', 'model'),
+        ('foo', 'data'))
+    with partitioning.axis_rules(p_rules):
+      self.assertEqual(
+          partitioning.logical_to_mesh_axes(('baz', 'bar', 'foo')),
+          ('data', None, 'model'))
 
   def test_logical_to_mesh_axes_priorities(self):
     p_rules = (