class_docstr: Extension of IndepMessenger that includes subsampling.
class_docstr_tok: ["Extension","of","IndepMessenger","that","includes","subsampling","."]
code: def _postprocess_message(self, msg: Message) -> None:
      if msg["type"] in ("param", "subsample") and self.dim is not None:
          event_dim = msg["kwargs"].get("event_dim")
          if event_dim is not None:
              assert event_dim >= 0
              dim = self.dim - event_dim
              assert msg["value"] is not None
              shape = msg["value"].shape
              if len(shape) >= -dim and shape[dim] != 1:
                  if (
                      is_validation_enabled()
                      and shape[dim] != self.size
                  ):
                      if msg["type"] == "param":
                          statement = "pyro.param({}, ..., event_dim={})".format(
                              msg["name"], event_dim
                          )
                      else:
                          statement = "pyro.subsample(..., event_dim={})".format(
                              event_dim
                          )
                      raise ValueError(
                          "Inside pyro.plate({}, {}, dim={}) invalid shape of {}: {}".format(
                              self.name,
                              self.size,
                              self.dim,
                              statement,
                              shape,
                          )
                      )
                  # Subsample parameters with known batch semantics.
                  if self.subsample_size < self.size:
                      value = msg["value"]
                      assert self._indices is not None
                      new_value = value.index_select(
                          dim, self._indices.to(value.device)
                      )
                      if msg["type"] == "param":
                          if hasattr(
                              value, "_pyro_unconstrained_param"
                          ):
                              param = value._pyro_unconstrained_param  # type: ignore[attr-defined]
                          else:
                              param = value.unconstrained()  # type: ignore[attr-defined]
  
                          if not hasattr(param, "_pyro_subsample"):
                              param._pyro_subsample = {}
  
                          param._pyro_subsample[dim] = self._indices
                          new_value._pyro_unconstrained_param = param  # type: ignore[attr-defined]
                      msg["value"] = new_value
  
file_path: pyro/pyro/poutine/subsample_messenger.py
imports: from typing import Optional, Tuple
  import torch
  from pyro.distributions.distribution import Distribution
  from pyro.poutine.indep_messenger import CondIndepStackFrame, IndepMessenger
  from pyro.poutine.runtime import Message, apply_stack
  from pyro.poutine.util import is_validation_enabled
  from pyro.util import ignore_jit_warnings
label_desc: Use image node_id 5 for calling the SubsampleMessenger obj's underlying member method code with example usage: obj._postprocess_message(msg) without return types