code: def _gen_partial_strategy(mesh: DeviceMesh) -> PlacementStrategy:
      """Util function to generate a partial strategy."""
      # NOTE: we use AVG by default, avg reduction is needed depending on
      # the loss function, for most loss function it should do
      # gradient averaging. There might be certain cases it should
      # not do gradient averaging (i.e. sum) but it's pretty rare.
      # TODO: Only NCCL supports AVG so using backend like Gloo would
      # crash, we should figure out a way to support avg reduction
      # for non-NCCL backend
      reduce_op = c10d.ReduceOp.AVG  # type: ignore[attr-defined]
      return PlacementStrategy(
          output_spec=DTensorSpec(
              mesh=mesh, placements=(_Partial(reduce_op),)
          ),
      )
  
imports: import operator
  from contextlib import contextmanager
  from enum import Enum
  from typing import Any, cast, Dict, List, Optional, Tuple
  import torch
  import torch.distributed.distributed_c10d
  import torch.fx
  import torch.library
  import torch.nn
  import torch.utils._pytree
  from torch.distributed._spmd.batch_dim_utils import BatchDimAnalyzer
  from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard
  from torch.distributed._tensor._utils import compute_local_shape
  from torch.distributed._tensor.op_schema import OpStrategy, PlacementStrategy, StrategyType, TupleStrategy
  from torch.distributed._tensor.placement_types import _Partial, DTensorSpec, Placement
  from torch.distributed._tensor.redistribute import redistribute_local_tensor
  from torch.fx import GraphModule
  from torch.fx.experimental.proxy_tensor import make_fx
  from torch.fx.passes.shape_prop import _extract_tensor_metadata
  from torch.nn.utils._named_member_accessor import NamedMemberAccessor
label_desc: Use image node_id 4 for calling a global function with example usage: _gen_partial_strategy(mesh) and returns: PlacementStrategy