code: def with_logical_partitioning(
      fn: Callable[..., Any],
      names: meta.LogicalNames,
      mesh: Optional[jax.sharding.Mesh] = None,
      rules: Optional[LogicalRules] = None,
  ) -> Callable[..., LogicallyPartitioned]:
      """Wraps a function's return value with LogicallyPartitioned.
  
      Example::
  
        >>> import flax.linen as nn
        >>> kernel_init = nn.with_logical_partitioning(
        ...     nn.initializers.lecun_normal, (None, "data"))
        >>> partitioned_dense = nn.Dense(features=3, kernel_init=kernel_init)
  
      Args:
        fn: The function to be wrapped. Typically this is an initializer.
        names: The logical axis passed to ``LogicallyPartitioned``.
        mesh: The mesh to use for the partitioning. If None, the global mesh
          resource is used if available.
        rules: Optional logical to mesh rules use. If None, the global rules
          are used if available.
      Returns:
        A function wrapping ``fn`` that will return an instance of
        ``LogicallyPartitioned``.
      """
  
      @functools.wraps(fn)
      def wrapper(*args, **kwargs):
          return LogicallyPartitioned(
              fn(*args, **kwargs), names, mesh=mesh, rules=rules
          )  # pytype: disable=wrong-keyword-args
  
      return wrapper
  
imports: import collections
  import contextlib
  import dataclasses
  import enum
  import functools
  import threading
  from typing import Any, Callable, List, Optional, Sequence, Tuple, Union
  import jax
  from jax import lax
  from jax.experimental import maps
  from flax import struct
  from flax.core import meta
label_desc: Use image node_id 14 for calling a global function with example usage: with_logical_partitioning(fn, names, mesh, rules) and returns: wrapper, LogicallyPartitioned