code: def param_with_axes(
      name: str,
      init_fn,
      *init_args,
      axes: Optional[Tuple[str, ...]] = None,
      module: Optional["nn.Module"] = None,
  ):
      """Declares and returns a parameter with logical axes in the current Module.
  
      See :mod:`flax.linen.module.param` for original docstring.
  
      Args:
        name: The parameter name.
        init_fn: The function that will be called to compute the initial value
          of this variable. This function will only be called the first time
          this parameter is used in this module.
        *init_args: The arguments to pass to init_fn.
        axes: A tuple of axis names, must match the rank of the param array.
        module: Use an explicit module instead of deriving the most recent from
          dynamic module context.
  
      Returns:
        The value of the initialized parameter.
  
      Raises:
        TypeError: if axes specification is mal-formed.
        ValueError: if specified logical axes don't match parameter rank.
      """
      # get current module if not explicitly provided
      if module is None:
          module = nn.module._context.module_stack[
              -1
          ]  # pylint: disable=protected-access
          assert module is not None
      # define/fetch parameter on that module
      module_param = module.param(name, init_fn, *init_args)
      if axes is not None:
          # apply logical axis constraint immediately
          module_param = with_sharding_constraint(
              module_param, jax.sharding.PartitionSpec(*axes)
          )
          # record logical axis constraint for global axis metadata
          module.sow(
              "params_axes",
              f"{name}_axes",
              AxisMetadata(axes),  # type: ignore
              reduce_fn=_param_with_axes_sow_reduce_fn,
          )
      return module_param
  
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 2 for calling a global function with example usage: param_with_axes(name, init_fn) and returns: module_param