text: import contextvars
  import logging
  import numbers
  import warnings
  
  from collections import UserDict
  from contextlib import AbstractContextManager
  from typing import (
      TYPE_CHECKING,
      Any,
      Callable,
      Dict,
      List,
      Optional,
      Set,
      Tuple,
      Union,
      cast,
      overload,
  )
  
  import numpy as np
  import theano
  import theano.tensor as tt
  
  from arviz import InferenceData
  from typing_extensions import Literal, Protocol
  from xarray import Dataset
  
  from pymc3.backends.base import MultiTrace
  from pymc3.distributions.distribution import (
      _compile_theano_function,
      _DrawValuesContext,
      _DrawValuesContextBlocker,
      is_fast_drawable,
      vectorized_ppc,
  )
  from pymc3.exceptions import IncorrectArgumentsError
  from pymc3.model import (
      Model,
      MultiObservedRV,
      ObservedRV,
      get_named_nodes_and_relations,
      modelcontext,
  )
  from pymc3.util import chains_and_samples, dataset_to_point_list, get_var_name
  from pymc3.vartypes import theano_constant
  
  # Failing tests:
  #    test_mixture_random_shape::test_mixture_random_shape
  #
  
  PosteriorPredictiveTrace = Dict[str, np.ndarray]
  Point = Dict[str, np.ndarray]
  
  
  class HasName(Protocol):
      name: str
  
  
  if TYPE_CHECKING:
      _TraceDictParent = UserDict[str, np.ndarray]
  else:
      _TraceDictParent = UserDict
  
  
  class _TraceDict(_TraceDictParent):
      """This class extends the standard trace-based representation
      of traces by adding some helpful attributes used in posterior predictive
      sampling.
  
      Attributes
      ~~~~~~~~~~
          varnames: list of strings"""
  
      varnames: List[str]
      _len: int
      data: Dict[str, np.ndarray]
  
      def __init__(
          self,
          point_list: Optional[List[Dict[str, np.ndarray]]] = None,
          multi_trace: Optional[MultiTrace] = None,
          dict: Optional[Dict[str, np.ndarray]] = None,
      ):
          """"""
          if multi_trace:
              assert point_list is None and dict is None
              self.data = {}  # Dict[str, np.ndarray]
              self._len = sum(len(multi_trace._straces[chain]) for chain in multi_trace.chains)
              self.varnames = multi_trace.varnames
              for vn in multi_trace.varnames:
                  self.data[vn] = multi_trace.get_values(vn)
          if point_list is not None:
              assert multi_trace is None and dict is None
              self.varnames = varnames = list(point_list[0].keys())
              rep_values = [point_list[0][varname] for varname in varnames]
              # translate the point list.
              self._len = num_points = len(point_list)
  
              def arr_for(val):
                  if np.isscalar(val):
                      return np.ndarray(shape=(num_points,))
                  elif isinstance(val, np.ndarray):
                      shp = (num_points,) + val.shape
                      return np.ndarray(shape=shp)
                  else:
                      raise TypeError(
                          "Illegal object %s of type %s as value of variable in point list."
                          % (val, type(val))
                      )
  
              self.data = {name: arr_for(val) for name, val in zip(varnames, rep_values)}
              for i, point in enumerate(point_list):
                  for var, value in point.items():
                      self.data[var][i] = value
          if dict is not None:
              assert point_list is None and multi_trace is None
              self.data = dict
              self.varnames = list(dict.keys())
              self._len = dict[self.varnames[0]].shape[0]
          assert self.varnames is not None and self._len is not None and self.data is not None
  
      def __len__(self) -> int:
          return self._len
  
      def _extract_slice(self, slc: slice) -> "_TraceDict":
          sliced_dict: Dict[str, np.ndarray] = {}
  
          def apply_slice(arr: np.ndarray) -> np.ndarray:
              if len(arr.shape) == 1:
                  return arr[slc]
              else:
                  return arr[slc, :]
  
          for vn, arr in self.data.items():
              sliced_dict[vn] = apply_slice(arr)
          return _TraceDict(dict=sliced_dict)
  
      @overload
      def __getitem__(self, item: Union[str, HasName]) -> np.ndarray:
          ...
  
      @overload
      def __getitem__(self, item: Union[slice, int]) -> "_TraceDict":
          ...
  
      def __getitem__(self, item):
          if isinstance(item, str):
              return super().__getitem__(item)
          elif isinstance(item, slice):
              return self._extract_slice(item)
          elif isinstance(item, int):
              return _TraceDict(dict={k: np.atleast_1d(v[item]) for k, v in self.data.items()})
          elif hasattr(item, "name"):
              return super().__getitem__(item.name)
          else:
              raise IndexError("Illegal index %s for _TraceDict" % str(item))
  
  
  def fast_sample_posterior_predictive(
      trace: Union[MultiTrace, Dataset, InferenceData, List[Dict[str, np.ndarray]]],
      samples: Optional[int] = None,
      model: Optional[Model] = None,
      var_names: Optional[List[str]] = None,
      keep_size: bool = False,
      random_seed=None,
  ) -> Dict[str, np.ndarray]:
      """Generate posterior predictive samples from a model given a trace.
  
      This is a vectorized alternative to the standard ``sample_posterior_predictive`` function.
      It aims to be as compatible as possible with the original API, and is significantly
      faster.  Both posterior predictive sampling functions have some remaining issues, and
      we encourage users to verify agreement across the results of both functions for the time
      being.
  
      Parameters
      ----------
      trace: MultiTrace, xarray.Dataset, InferenceData, or List of points (dictionary)
          Trace generated from MCMC sampling.
      samples: int, optional
          Number of posterior predictive samples to generate. Defaults to one posterior predictive
          sample per posterior sample, that is, the number of draws times the number of chains. It
          is not recommended to modify this value; when modified, some chains may not be represented
          in the posterior predictive sample.
      model: Model (optional if in `with` context)
          Model used to generate `trace`
      var_names: Iterable[str]
          List of vars to sample.
      keep_size: bool, optional
          Force posterior predictive sample to have the same shape as posterior and sample stats
          data: ``(nchains, ndraws, ...)``.
      random_seed: int
          Seed for the random number generator.
  
      Returns
      -------
      samples: dict
          Dictionary with the variable names as keys, and values numpy arrays containing
          posterior predictive samples.
      """
  
      ### Implementation note: primarily this function canonicalizes the arguments:
      ### Establishing the model context, wrangling the number of samples,
      ### Canonicalizing the trace argument into a _TraceDict object and fitting it
      ### to the requested number of samples.  Then it invokes posterior_predictive_draw_values
      ### *repeatedly*.  It does this repeatedly, because the trace argument is set up to be
      ### the same as the number of samples. So if the number of samples requested is
      ### greater than the number of samples in the trace parameter, we sample repeatedly.  This
      ### makes the shape issues just a little easier to deal with.
  
      if isinstance(trace, InferenceData):
          nchains, ndraws = chains_and_samples(trace)
          trace = dataset_to_point_list(trace.posterior)
      elif isinstance(trace, Dataset):
          nchains, ndraws = chains_and_samples(trace)
          trace = dataset_to_point_list(trace)
      elif isinstance(trace, MultiTrace):
          nchains = trace.nchains
          ndraws = len(trace)
      else:
          if keep_size:
              # arguably this should be just a warning.
              raise IncorrectArgumentsError(
                  "For keep_size, cannot identify chains and length from %s.", trace
              )
  
      model = modelcontext(model)
      assert model is not None
      with model:
  
          if keep_size and samples is not None:
              raise IncorrectArgumentsError("Should not specify both keep_size and samples arguments")
  
          if isinstance(trace, list) and all(isinstance(x, dict) for x in trace):
              _trace = _TraceDict(point_list=trace)
          elif isinstance(trace, MultiTrace):
              _trace = _TraceDict(multi_trace=trace)
          else:
              raise TypeError(
                  "Unable to generate posterior predictive samples from argument of type %s"
                  % type(trace)
              )
  
          len_trace = len(_trace)
  
          assert isinstance(_trace, _TraceDict)
  
          _samples: List[int] = []
          # temporary replacement for more complicated logic.
          max_samples: int = len_trace
          if samples is None or samples == max_samples:
              _samples = [max_samples]
          elif samples < max_samples:
              warnings.warn(
                  "samples parameter is smaller than nchains times ndraws, some draws "
                  "and/or chains may not be represented in the returned posterior "
                  "predictive sample"
              )
              # if this is less than the number of samples in the trace, take a slice and
              # work with that.
              _trace = _trace[slice(samples)]
              _samples = [samples]
          elif samples > max_samples:
              full, rem = divmod(samples, max_samples)
              _samples = (full * [max_samples]) + ([rem] if rem != 0 else [])
          else:
              raise IncorrectArgumentsError(
                  "Unexpected combination of samples (%s) and max_samples (%d)"
                  % (samples, max_samples)
              )
  
          if var_names is None:
              vars = model.observed_RVs
          else:
              vars = [model[x] for x in var_names]
  
          if random_seed is not None:
              np.random.seed(random_seed)
  
          if TYPE_CHECKING:
              _ETPParent = UserDict[str, np.ndarray]  # this is only processed by mypy
          else:
              # this is not seen by mypy but will be executed at runtime.
              _ETPParent = UserDict
  
          class _ExtendableTrace(_ETPParent):
              def extend_trace(self, trace: Dict[str, np.ndarray]) -> None:
                  for k, v in trace.items():
                      if k in self.data:
                          self.data[k] = np.concatenate((self.data[k], v))
                      else:
                          self.data[k] = v
  
          ppc_trace = _ExtendableTrace()
          for s in _samples:
              strace = _trace if s == len_trace else _trace[slice(0, s)]
              try:
                  values = posterior_predictive_draw_values(cast(List[Any], vars), strace, s)
                  new_trace: Dict[str, np.ndarray] = {k.name: v for (k, v) in zip(vars, values)}
                  ppc_trace.extend_trace(new_trace)
              except KeyboardInterrupt:
                  pass
  
          if keep_size:
              return {
                  k: ary.reshape((nchains, ndraws, *ary.shape[1:])) for k, ary in ppc_trace.items()
              }
          # this gets us a Dict[str, np.ndarray] instead of my wrapped equiv.
          return ppc_trace.data
  
  
  def posterior_predictive_draw_values(
      vars: List[Any], trace: _TraceDict, samples: int
  ) -> List[np.ndarray]:
      with _PosteriorPredictiveSampler(vars, trace, samples, None) as sampler:
          return sampler.draw_values()
  
  
  class _PosteriorPredictiveSampler(AbstractContextManager):
      """The process of posterior predictive sampling is quite complicated so this provides a central data store."""
  
      # inputs
      vars: List[Any]
      trace: _TraceDict
      samples: int
      size: Optional[int]  # not supported!
  
      # other slots
      logger: logging.Logger
  
      # for the search
      evaluated: Dict[int, np.ndarray]
      symbolic_params: List[Tuple[int, Any]]
  
      # set by make_graph...
      leaf_nodes: Dict[str, Any]
      named_nodes_parents: Dict[str, Any]
      named_nodes_children: Dict[str, Any]
      _tok: contextvars.Token
  
      def __init__(self, vars, trace: _TraceDict, samples, model: Optional[Model], size=None):
          if size is not None:
              raise NotImplementedError(
                  "sample_posterior_predictive does not support the size argument at this time."
              )
          assert vars is not None
          self.vars = vars
          self.trace = trace
          self.samples = samples
          self.size = size
          self.logger = logging.getLogger("posterior_predictive")
  
      def __enter__(self) -> "_PosteriorPredictiveSampler":
          self._tok = vectorized_ppc.set(posterior_predictive_draw_values)
          return self
  
      def __exit__(self, exc_type, exc_val, exc_tb) -> Literal[False]:
          vectorized_ppc.reset(self._tok)
          return False
  
      def draw_values(self) -> List[np.ndarray]:
          vars = self.vars
          trace = self.trace
          samples = self.samples
          # size = self.size
          params = dict(enumerate(vars))
  
          with _DrawValuesContext() as context:
              self.init()
              self.make_graph()
  
              drawn = context.drawn_vars
  
              # Init givens and the stack of nodes to try to `_draw_value` from
              givens = {
                  p.name: (p, v)
                  for (p, samples), v in drawn.items()
                  if getattr(p, "name", None) is not None
              }
              stack = list(self.leaf_nodes.values())  # A queue would be more appropriate
  
              while stack:
                  next_ = stack.pop(0)
                  if (next_, samples) in drawn:
                      # If the node already has a givens value, skip it
                      continue
                  elif isinstance(next_, (theano_constant, tt.sharedvar.SharedVariable)):
                      # If the node is a theano.tensor.TensorConstant or a
                      # theano.tensor.sharedvar.SharedVariable, its value will be
                      # available automatically in _compile_theano_function so
                      # we can skip it. Furthermore, if this node was treated as a
                      # TensorVariable that should be compiled by theano in
                      # _compile_theano_function, it would raise a `TypeError:
                      # ('Constants not allowed in param list', ...)` for
                      # TensorConstant, and a `TypeError: Cannot use a shared
                      # variable (...) as explicit input` for SharedVariable.
                      # ObservedRV and MultiObservedRV instances are ViewOPs
                      # of TensorConstants or SharedVariables, we must add them
                      # to the stack or risk evaluating deterministics with the
                      # wrong values (issue #3354)
                      stack.extend(
                          [
                              node
                              for node in self.named_nodes_parents[next_]
                              if isinstance(node, (ObservedRV, MultiObservedRV))
                              and (node, samples) not in drawn
                          ]
                      )
                      continue
                  else:
                      # If the node does not have a givens value, try to draw it.
                      # The named node's children givens values must also be taken
                      # into account.
                      children = self.named_nodes_children[next_]
                      temp_givens = [givens[k] for k in givens if k in children]
                      try:
                          # This may fail for autotransformed RVs, which don't
                          # have the random method
                          value = self.draw_value(next_, trace=trace, givens=temp_givens)
                          assert isinstance(value, np.ndarray)
                          givens[next_.name] = (next_, value)
                          drawn[(next_, samples)] = value
                      except theano.gof.fg.MissingInputError:
                          # The node failed, so we must add the node's parents to
                          # the stack of nodes to try to draw from. We exclude the
                          # nodes in the `params` list.
                          stack.extend(
                              [
                                  node
                                  for node in self.named_nodes_parents[next_]
                                  if node is not None and (node, samples) not in drawn
                              ]
                          )
  
              # the below makes sure the graph is evaluated in order
              # test_distributions_random::TestDrawValues::test_draw_order fails without it
              # The remaining params that must be drawn are all hashable
              to_eval: Set[int] = set()
              missing_inputs: Set[int] = {j for j, p in self.symbolic_params}
  
              while to_eval or missing_inputs:
                  if to_eval == missing_inputs:
                      raise ValueError(
                          "Cannot resolve inputs for {}".format(
                              [get_var_name(trace.varnames[j]) for j in to_eval]
                          )
                      )
                  to_eval = set(missing_inputs)
                  missing_inputs = set()
                  for param_idx in to_eval:
                      param = vars[param_idx]
                      drawn = context.drawn_vars
                      if (param, samples) in drawn:
                          self.evaluated[param_idx] = drawn[(param, samples)]
                      else:
                          try:
                              if param in self.named_nodes_children:
                                  for node in self.named_nodes_children[param]:
                                      if node.name not in givens and (node, samples) in drawn:
                                          givens[node.name] = (
                                              node,
                                              drawn[(node, samples)],
                                          )
                              value = self.draw_value(param, trace=self.trace, givens=givens.values())
                              assert isinstance(value, np.ndarray)
                              self.evaluated[param_idx] = drawn[(param, samples)] = value
                              givens[param.name] = (param, value)
                          except theano.gof.fg.MissingInputError:
                              missing_inputs.add(param_idx)
          return [self.evaluated[j] for j in params]
  
      def init(self) -> None:
          """This method carries out the initialization phase of sampling
          from the posterior predictive distribution.  Notably it initializes the
          ``_DrawValuesContext`` bookkeeping object and evaluates the "fast drawable"
          parts of the model."""
          vars: List[Any] = self.vars
          trace: _TraceDict = self.trace
          samples: int = self.samples
          leaf_nodes: Dict[str, Any]
          named_nodes_parents: Dict[str, Any]
          named_nodes_children: Dict[str, Any]
  
          # initialization phase
          context = _DrawValuesContext.get_context()
          assert isinstance(context, _DrawValuesContext)
          with context:
              drawn = context.drawn_vars
              evaluated: Dict[int, Any] = {}
              symbolic_params = []
              for i, var in enumerate(vars):
                  if is_fast_drawable(var):
                      evaluated[i] = self.draw_value(var)
                      continue
                  name = getattr(var, "name", None)
                  if (var, samples) in drawn:
                      evaluated[i] = drawn[(var, samples)]
                      # We filter out Deterministics by checking for `model` attribute
                  elif name is not None and hasattr(var, "model") and name in trace.varnames:
                      # param.name is in the trace.  Record it as drawn and evaluated
                      drawn[(var, samples)] = evaluated[i] = trace[cast(str, name)]
                  else:
                      # param still needs to be drawn
                      symbolic_params.append((i, var))
          self.evaluated = evaluated
          self.symbolic_params = symbolic_params
  
      def make_graph(self) -> None:
          # Distribution parameters may be nodes which have named node-inputs
          # specified in the point. Need to find the node-inputs, their
          # parents and children to replace them.
          symbolic_params = self.symbolic_params
          self.leaf_nodes = {}
          self.named_nodes_parents = {}
          self.named_nodes_children = {}
          for _, param in symbolic_params:
              if hasattr(param, "name"):
                  # Get the named nodes under the `param` node
                  nn, nnp, nnc = get_named_nodes_and_relations(param)
                  self.leaf_nodes.update(nn)
                  # Update the discovered parental relationships
                  for k in nnp.keys():
                      if k not in self.named_nodes_parents.keys():
                          self.named_nodes_parents[k] = nnp[k]
                      else:
                          self.named_nodes_parents[k].update(nnp[k])
                  # Update the discovered child relationships
                  for k in nnc.keys():
                      if k not in self.named_nodes_children.keys():
                          self.named_nodes_children[k] = nnc[k]
                      else:
                          self.named_nodes_children[k].update(nnc[k])
  
      def draw_value(self, param, trace: Optional[_TraceDict] = None, givens=None):
          """Draw a set of random values from a distribution or return a constant.
  
          Parameters
          ----------
          param: number, array like, theano variable or pymc3 random variable
              The value or distribution. Constants or shared variables
              will be converted to an array and returned. Theano variables
              are evaluated. If `param` is a pymc3 random variable, draw
              values from it and return that (as ``np.ndarray``), unless a
              value is specified in the ``trace``.
          trace: pm.MultiTrace, optional
              A dictionary from pymc3 variable names to samples of their values
              used to provide context for evaluating ``param``.
          givens: dict, optional
              A dictionary from theano variables to their values. These values
              are used to evaluate ``param`` if it is a theano variable.
          """
          samples = self.samples
  
          def random_sample(
              meth: Callable[..., np.ndarray],
              param,
              point: _TraceDict,
              size: int,
              shape: Tuple[int, ...],
          ) -> np.ndarray:
              val = meth(point=point, size=size)
              try:
                  assert val.shape == (size,) + shape, (
                      "Sampling from random of %s yields wrong shape" % param
                  )
              # error-quashing here is *extremely* ugly, but it seems to be what the logic in DensityDist wants.
              except AssertionError as e:
                  if (
                      hasattr(param, "distribution")
                      and hasattr(param.distribution, "wrap_random_with_dist_shape")
                      and not param.distribution.wrap_random_with_dist_shape
                  ):
                      pass
                  else:
                      raise e
  
              return val
  
          if isinstance(param, (numbers.Number, np.ndarray)):
              return param
          elif isinstance(param, theano_constant):
              return param.value
          elif isinstance(param, tt.sharedvar.SharedVariable):
              return param.get_value()
          elif isinstance(param, (tt.TensorVariable, MultiObservedRV)):
              if hasattr(param, "model") and trace and param.name in trace.varnames:
                  return trace[param.name]
              elif hasattr(param, "random") and param.random is not None:
                  model = modelcontext(None)
                  assert isinstance(model, Model)
                  shape: Tuple[int, ...] = tuple(_param_shape(param, model))
                  return random_sample(param.random, param, point=trace, size=samples, shape=shape)
              elif (
                  hasattr(param, "distribution")
                  and hasattr(param.distribution, "random")
                  and param.distribution.random is not None
              ):
                  if hasattr(param, "observations"):
                      # shape inspection for ObservedRV
                      dist_tmp = param.distribution
                      try:
                          distshape: Tuple[int, ...] = tuple(param.observations.shape.eval())
                      except AttributeError:
                          distshape = tuple(param.observations.shape)
  
                      dist_tmp.shape = distshape
                      try:
                          return random_sample(
                              dist_tmp.random,
                              param,
                              point=trace,
                              size=samples,
                              shape=distshape,
                          )
                      except (ValueError, TypeError):
                          # reset shape to account for shape changes
                          # with theano.shared inputs
                          dist_tmp.shape = ()
                          # We want to draw values to infer the dist_shape,
                          # we don't want to store these drawn values to the context
                          with _DrawValuesContextBlocker():
                              point = trace[0] if trace else None
                              temp_val = np.atleast_1d(dist_tmp.random(point=point, size=None))
                          # if hasattr(param, 'name') and param.name == 'obs':
                          #     import pdb; pdb.set_trace()
                          # Sometimes point may change the size of val but not the
                          # distribution's shape
                          if point and samples is not None:
                              temp_size = np.atleast_1d(samples)
                              if all(temp_val.shape[: len(temp_size)] == temp_size):
                                  dist_tmp.shape = tuple(temp_val.shape[len(temp_size) :])
                              else:
                                  dist_tmp.shape = tuple(temp_val.shape)
                          # I am not sure why I need to do this, but I do in order to trim off a
                          # degenerate dimension [2019/09/05:rpg]
                          if dist_tmp.shape[0] == 1 and len(dist_tmp.shape) > 1:
                              dist_tmp.shape = dist_tmp.shape[1:]
                          return random_sample(
                              dist_tmp.random,
                              point=trace,
                              size=samples,
                              param=param,
                              shape=tuple(dist_tmp.shape),
                          )
                  else:  # has a distribution, but no observations
                      distshape = tuple(param.distribution.shape)
                      return random_sample(
                          meth=param.distribution.random,
                          param=param,
                          point=trace,
                          size=samples,
                          shape=distshape,
                      )
              # NOTE: I think the following is already vectorized.
              else:
                  if givens:
                      variables, values = list(zip(*givens))
                  else:
                      variables = values = []
                  # We only truly care if the ancestors of param that were given
                  # value have the matching dshape and val.shape
                  param_ancestors = set(theano.gof.graph.ancestors([param], blockers=list(variables)))
                  inputs = [
                      (var, val) for var, val in zip(variables, values) if var in param_ancestors
                  ]
                  if inputs:
                      input_vars, input_vals = list(zip(*inputs))
                  else:
                      input_vars = []
                      input_vals = []
                  func = _compile_theano_function(param, input_vars)
                  if not input_vars:
                      assert input_vals == []  # AFAICT if there are now vars, there can't be vals
                      output = func(*input_vals)
                      if hasattr(output, "shape"):
                          val = np.repeat(np.expand_dims(output, 0), samples, axis=0)
                      else:
                          val = np.full(samples, output)
  
                  else:
                      val = func(*input_vals)
                      # np.ndarray([func(*input_vals) for inp in zip(*input_vals)])
                  return val
          raise ValueError("Unexpected type in draw_value: %s" % type(param))
  
  
  def _param_shape(var_desig, model: Model) -> Tuple[int, ...]:
      if isinstance(var_desig, str):
          v = model[var_desig]
      else:
          v = var_desig
      if hasattr(v, "observations"):
          try:
              # To get shape of _observed_ data container `pm.Data`
              # (wrapper for theano.SharedVariable) we evaluate it.
              shape = tuple(v.observations.shape.eval())
          except AttributeError:
              shape = v.observations.shape
      elif hasattr(v, "dshape"):
          shape = v.dshape
      else:
          shape = v.tag.test_value.shape
      if shape == (1,):
          shape = tuple()
      return shape
  