code: def with_partitioning(
      initializer: F,
      sharding: Sharding,
      mesh: tp.Optional[jax.sharding.Mesh] = None,
      get_value_hooks: tp.Union[
          variables.GetValueHook[A],
          tp.Sequence[variables.GetValueHook[A]],
      ] = (),
      create_value_hooks: tp.Union[
          variables.CreateValueHook[A],
          tp.Sequence[variables.CreateValueHook[A]],
      ] = (),
      **metadata: tp.Any,
  ) -> F:
      if callable(get_value_hooks):
          get_value_hooks = (get_value_hooks, sharding_hook)
      else:
          get_value_hooks = (*get_value_hooks, sharding_hook)
  
      if callable(create_value_hooks):
          create_value_hooks = (create_value_hooks, sharding_hook)
      else:
          create_value_hooks = (*create_value_hooks, sharding_hook)
  
      return variables.with_metadata(
          initializer,
          get_value_hooks=get_value_hooks,
          create_value_hooks=create_value_hooks,
          sharding=sharding,
          mesh=mesh,
          **metadata,
      )
  
imports: import functools
  import typing
  import jax
  from jax.experimental import maps
  from jax.sharding import Mesh, PartitionSpec
  from flax.experimental.nnx.nnx import variables
  from flax.experimental.nnx.nnx.pytreelib import TreeNode
  from flax.experimental.nnx.nnx.state import State
label_desc: Use image node_id 11 for calling a global function with example usage: with_partitioning(initializer, sharding, mesh, get_value_hooks, create_value_hooks) and returns: variables