class_docstr: Enumerates in parallel over discrete sample sites marked
  ``infer={"enumerate": "parallel"}``.
  
  :param int first_available_dim: The first tensor dimension (counting
      from the right) that is available for parallel enumeration. This
      dimension and all dimensions left may be used internally by Pyro.
      This should be a negative integer or None.
class_docstr_tok: ["Enumerates","in","parallel","over","discrete","sample","sites","marked","``","infer=","{","``","enumerate","''",":","``","parallel","''","}","``",".",":","param","int","first_available_dim",":","The","first","tensor","dimension","(","counting","from","the","right",")","that","is","available","for","parallel","enumeration",".","This","dimension","and","all","dimensions","left","may","be","used","internally","by","Pyro",".","This","should","be","a","negative","integer","or","None","."]
docstr: :param msg: current message at a trace site.
  :returns: a sample from the stochastic function at the site.
docstr_tok: [":","param","msg",":","current","message","at","a","trace","site",".",":","returns",":","a","sample","from","the","stochastic","function","at","the","site","."]
code: def _pyro_sample(self, msg: Message) -> None:
      """
      :param msg: current message at a trace site.
      :returns: a sample from the stochastic function at the site.
      """
      if msg["done"] or not isinstance(
          msg["fn"], TorchDistributionMixin
      ):
          return
  
      assert isinstance(msg["name"], str)
      assert msg["infer"] is not None
      # Compute upstream dims in scope; these are unsafe to use for this site's target_dim.
      scope = msg["infer"].get(
          "_markov_scope"
      )  # site name -> markov depth
      param_dims = (
          _ENUM_ALLOCATOR.dim_to_id.copy()
      )  # enum dim -> unique id
      if scope is not None:
          for name, depth in scope.items():
              if (
                  self._markov_depths[name] == depth
              ):  # hide sites whose markov context has exited
                  param_dims.update(self._value_dims[name])
          self._markov_depths[msg["name"]] = msg["infer"][
              "_markov_depth"
          ]
      self._param_dims[msg["name"]] = param_dims
      if (
          msg["is_observed"]
          or msg["infer"].get("enumerate") != "parallel"
      ):
          return
  
      # Compute an enumerated value (at an arbitrary dim).
      value = enumerate_site(msg)
      actual_dim = -1 - len(
          msg["fn"].batch_shape
      )  # the leftmost dim of log_prob
  
      # Move actual_dim to a safe target_dim.
      target_dim, id_ = _ENUM_ALLOCATOR.allocate(
          None if scope is None else set(param_dims)
      )
      event_dim = msg["fn"].event_dim
      categorical_support = getattr(
          value, "_pyro_categorical_support", None
      )
      if categorical_support is not None:
          # Preserve categorical supports to speed up Categorical.log_prob().
          # See pyro/distributions/torch.py for details.
          assert target_dim < 0
          value = value.reshape(
              value.shape[:1] + (1,) * (-1 - target_dim)
          )
          value._pyro_categorical_support = categorical_support  # type: ignore[attr-defined]
      elif actual_dim < target_dim:
          assert (
              value.size(target_dim - event_dim) == 1
          ), "pyro.markov dim conflict at dim {}".format(actual_dim)
          value = value.transpose(
              target_dim - event_dim, actual_dim - event_dim
          )
          while value.dim() and value.size(0) == 1:
              value = value.squeeze(0)
      elif target_dim < actual_dim:
          diff = actual_dim - target_dim
          value = value.reshape(
              value.shape[:1] + (1,) * diff + value.shape[1:]
          )
  
      # Compute dims passed downstream through the value.
      value_dims = {
          dim: param_dims[dim]
          for dim in range(event_dim - value.dim(), 0)
          if value.size(dim - event_dim) > 1 and dim in param_dims
      }
      value_dims[target_dim] = id_
  
      msg["infer"]["_enumerate_dim"] = target_dim
      msg["infer"]["_dim_to_id"] = value_dims
      msg["value"] = value
      msg["done"] = True
  
file_path: pyro/pyro/poutine/enum_messenger.py
imports: from typing import Any, Dict, List, Optional
  import torch
  from typing_extensions import Self
  from pyro.distributions.torch import Categorical
  from pyro.distributions.torch_distribution import TorchDistributionMixin
  from pyro.ops.indexing import Vindex
  from pyro.poutine.messenger import Messenger
  from pyro.poutine.runtime import _ENUM_ALLOCATOR, Message
  from pyro.util import ignore_jit_warnings
label_desc: Use image node_id 3 for calling the EnumMessenger obj's underlying member method code with example usage: obj._pyro_sample(msg) without return types