code: def with_partitioning(
      fn: Callable[..., Any],
      names: LogicalNames,
      mesh: Optional[jax.sharding.Mesh] = None,
  ) -> Callable[..., Partitioned[Any]]:
      """Wraps a function's return value with Partitioned.
  
      Example::
  
        kernel_init = with_partitioning(
            nn.initializers.lecun_normal, (None, "data"))
        partitioned_dense = nn.Dense(features, kernel_init=kernel_init)
  
      Args:
        fn: The function to be wrapped. Typically this is an initializer.
        names: The logical axis passed to ``Partitioned``.
        mesh: The mesh to use for the partitioning. If None, the global mesh
          resource is used if available.
      Returns:
        A function wrapping ``fn`` that will return an instance of ``Partitioned``.
      """
  
      @functools.wraps(fn)
      def wrapper(*args, **kwargs):
          return Partitioned(fn(*args, **kwargs), names, mesh=mesh)
  
      return wrapper
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Generic, Optional, Tuple, TypeVar, Union
  import jax
  from jax.experimental import maps
  from flax import errors, struct
label_desc: Use image node_id 8 for calling a global function with example usage: with_partitioning(fn, names, mesh) and returns: wrapper, Partitioned