How to pjit initial state of train_state object?

[mattiasmar]

Hello,
In this collab link the layers of the MNIST CNN is meant to be pjit'ed according to the simple rules:

"bias" -> P(None, None)
"kernel" -> P(None, "MP")

The pjit call is given as:

p_get_initial_state = pjit(
    get_initial_state,
    in_axis_resources=None,
    out_axis_resources=(opt_state_spec, param_spec),
)

The details of how opt_state_spec & param_spec are defined is given in the very last cell of the collab.

When calling the pjit'ed method in the scope of the mesh one gets the below error (indentation added by me for increased readability).
Q1: Is the error due to the "None" high-lighted in the below?
Q2: Can you see how I can overcome this error and create a valid out_axis_resources mapping?

ValueError: pjit out_axis_resources specification must be a tree prefix of the corresponding value, got specification 
((TraceState(trace=FrozenDict({
    Conv_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
    Dense_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
})), None),   #<-------Q1: Is the error due to this value?
 FrozenDict({
    Conv_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
    Dense_0: FrozenDict({
        bias: PartitionSpec(None, None),
        kernel: PartitionSpec(None, 'MP'),
    }),
})) 

for value tree PyTreeDef
(
  (
    (
      CustomNode(namedtuple[<class 'optax._src.transform.TraceState'>],
      [CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()],
      [{'Conv_0': {'bias': *, 'kernel': *}, 'Dense_0': {'bias': *, 'kernel': *}}])]
      ), 
      CustomNode(namedtuple[<class 'optax._src.base.EmptyState'>], [])
    ),
    CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], 
    [{'Conv_0': {'bias': *, 'kernel': *}, 'Dense_0': {'bias': *, 'kernel': *}}])
  )
)
  

Thank you!

[mattiasmar]

Answered in google/flax#1789