code: def __init__(
      self,
      propagate_user_sharding,
      partition,
      to_mesh_pspec_sharding,
      in_tree,
      out_tree,
      infer_sharding_from_operands,
      module_context,
      mesh,
      static_args,
  ):
      self.propagate_user_sharding = propagate_user_sharding
      self.partition = partition
      self.to_mesh_pspec_sharding = to_mesh_pspec_sharding
      self.in_tree = in_tree
      self.out_tree = out_tree
      self.infer_sharding_from_operands = infer_sharding_from_operands
      self.module_context = module_context
      self.mesh = mesh
      self.static_args = static_args
  
imports: from __future__ import annotations
  import inspect
  from typing import Optional
  import weakref
  import jax
  from jax._src import core
  from jax import tree_util
  from jax._src import linear_util
  from jax._src import sharding_impls
  from jax.errors import UnexpectedTracerError
  from jax._src import mesh
  from jax._src import dispatch
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.lib.mlir import ir
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.sharding_impls import _op_sharding_to_pos_sharding
  from jax._src import custom_api_util
  from jax._src.lib import xla_client
  from jax._src.api_util import flatten_fun_nokwargs, argnums_partial
label_desc: Use image node_id 1 to create a new _ShardingCallbackInfo object  with example: obj = _ShardingCallbackInfo(propagate_user_sharding, partition, to_mesh_pspec_sharding, in_tree, out_tree, infer_sharding_from_operands, module_context, mesh, static_args)