code: def _param_with_axes_sow_reduce_fn(x, y):
      """Reduction function for sow() calls.
  
      Args:
        x: Existing value, or () if there was none.
        y: New axis names sown.
  
      Returns:
        New axis names.
  
      Raises:
        TypeError: If the newly sown value is not an AxisMetadata.
        ValueError: If the newly sown axis names don't match previously sown axis
          names.
        AssertionError: If a previously sown value was truthy and not an
          AxisMetadata.
      """
      if not isinstance(y, AxisMetadata):
          raise TypeError(
              "Expected newly sown value to be an AxisMetadata"
          )
  
      if isinstance(x, AxisMetadata):
          if x != y:
              raise ValueError(
                  "If axis names are sown twice, expected them to match. "
                  f"Got {x} and {y}."
              )
      elif x:
          # Shouldn't happen, so raise a fairly internal error.
          raise AssertionError(
              f"Non-initial-or-AxisMetadata value encountered: {x}"
          )
      return y
  
imports: import functools
  import re
  from typing import Any, Callable, Mapping, Optional, Tuple
  import flax
  from flax import linen
  from flax import struct
  from flax.core.frozen_dict import freeze
  from flax.core.frozen_dict import unfreeze
  from flax.core.lift import In
  from flax.core.lift import Out
  from flax.linen.spmd import _axis_rules
  from flax.linen.spmd import _AxisRules
  from flax.linen.spmd import _is_logical_spec
  from flax.linen.spmd import _with_sharding_constraint
  from flax.linen.spmd import Array
  from flax.linen.spmd import ArrayPytree
  from flax.linen.spmd import get_logical_axis_rules
  from flax.linen.spmd import logical_axis_rules
  from flax.linen.spmd import logical_to_mesh
  from flax.linen.spmd import logical_to_mesh_axes
  from flax.linen.spmd import LogicalPartitionSpec
  from flax.linen.spmd import LogicalPartitionSpecPytree
  from flax.linen.spmd import LogicalRules
  from flax.linen.spmd import PartitionSpecPytree
  from flax.linen.spmd import RulesFallback
  from flax.linen.spmd import set_logical_axis_rules
  from flax.linen.spmd import with_logical_constraint
  from flax.traverse_util import flatten_dict
  from flax.traverse_util import unflatten_dict
  import jax
label_desc: Use image node_id 1 for calling a global function with example usage: _param_with_axes_sow_reduce_fn(x, y) and returns: y