tensorflow_probability/python/distributions/internal/slicing.py
def batch_slice(dist, params_event_ndims, params_overrides, slices):
  """Slices `dist` along its batch dimensions. Helper for tfd.Distribution.

  Args:
    dist: A `tfd.Distribution` instance.
    params_event_ndims: A `dict` of `str->int` indicating the number of
      dimensions of a given parameter required to parameterize a single event.
    params_overrides: A `dict` of parameter overrides. (e.g. from
      `Distribution.copy`).
    slices: A `slice` or `int` or `int` `Tensor` or `tf.newaxis` or `tuple`
      thereof. (e.g. the argument of a `__getitem__` method).

  Returns:
    new_dist: A batch-sliced `tfd.Distribution`.
  """
  if not isinstance(slices, collections.Sequence):
    slices = (slices,)
  # We track the history of slice and copy(**param_overrides) in order to trace
  # back to the original distribution's source variables.
  orig_dist, slice_overrides_seq = getattr(dist, PROVENANCE_ATTR, (dist, []))
  slice_overrides_seq += [(slices, params_overrides)]
  # Re-doing the full sequence of slice+copy override work here enables
  # gradients all the way back to the original distribution's arguments.
  dist = _apply_slice_sequence(orig_dist, params_event_ndims,
                               slice_overrides_seq)
  setattr(dist, PROVENANCE_ATTR, (orig_dist, slice_overrides_seq))
  return dist
def batch_slice(dist, params_event_ndims, params_overrides, slices):
  """Slices `dist` along its batch dimensions. Helper for tfd.Distribution.

  Args:
    dist: A `tfd.Distribution` instance.
    params_event_ndims: A `dict` of `str->int` indicating the number of
      dimensions of a given parameter required to parameterize a single event.
    params_overrides: A `dict` of parameter overrides. (e.g. from
      `Distribution.copy`).
    slices: A `slice` or `int` or `int` `Tensor` or `tf.newaxis` or `tuple`
      thereof. (e.g. the argument of a `__getitem__` method).

  Returns:
    new_dist: A batch-sliced `tfd.Distribution`.
  """
  if not isinstance(slices, collections.Sequence):
    slices = (slices,)
  # We track the history of slice and copy(**param_overrides) in order to trace
  # back to the original distribution's source variables.
  orig_dist, slice_overrides_seq = getattr(dist, PROVENANCE_ATTR, (dist, []))
  slice_overrides_seq += [(slices, params_overrides)]
  # Re-doing the full sequence of slice+copy override work here enables
  # gradients all the way back to the original distribution's arguments.
  dist = _apply_slice_sequence(orig_dist, params_event_ndims,
                               slice_overrides_seq)
  setattr(dist, PROVENANCE_ATTR, (orig_dist, slice_overrides_seq))
  return dist
Slices `dist` along its batch dimensions. Helper for tfd.Distribution.

  Args:
    dist: A `tfd.Distribution` instance.
    params_event_ndims: A `dict` of `str->int` indicating the number of
      dimensions of a given parameter required to parameterize a single event.
    params_overrides: A `dict` of parameter overrides. (e.g. from
      `Distribution.copy`).
    slices: A `slice` or `int` or `int` `Tensor` or `tf.newaxis` or `tuple`
      thereof. (e.g. the argument of a `__getitem__` method).

  Returns:
    new_dist: A batch-sliced `tfd.Distribution`.