tfp.experimental.unnest.get_nested_objs

Finds the list of nested objects inside an object's attributes.

get_nested_objs proceeds as follow:

  1. If hasattr(obj, nested_key), then set nested_attrs = getattr(obj, nested_key).
  2. Otherwise set nested_attrs = fallback_attrs.
  3. In either case, nested_attrs should now be a string or collection of strings. Return the list [(attr, getattr(obj, attr)) for attr in nested_attrs] omitting missing attributes.

nested_key is for class- or object-level customization, and fallback_attrs for invocation-level customization.

Example:

class Nest:
  __nested_attrs__ = ('inner1', 'inner2')

  def __init__(self, inner1, inner2, inner3):
    self.inner1 = inner1
    self.inner2 = inner2
    self.inner3 = inner3

nest = Nest('x', 'y', 'z')

# Search stops if `nested_key` is found.
get_nested_objs(
    nest,
    nested_key='__nested_attrs__',
    fallback_attrs=('inner1', 'inner2', 'inner3')) == [
        ('inner1', 'x'), ('inner2', 'y')]  # True

# If `nested_key` is not found, search in all of `fallback_attrs`.
get_nested_objs(
    nest,
    nested_key='does_not_exist',
    fallback_attrs=('inner1', 'inner2', 'inner3')) == [
        ('inner1', 'x'), ('inner2', 'y'), ('inner3', 'z')]  # True

# If nothing is found, empty list is returned.
get_nested_objs(
    nest,
    nested_key='does_not_exist',
    fallback_attrs=('does_not_exist')) == []  # True

# `getattr(obj, nested_key)` and `fallback_attrs` can be either strings or
# collections of strings.
nest2 = Nest('x', 'y', 'z')
nest2.__nested_attrs__ = 'inner3'
get_nested_objs(
    nest2,
    nested_key='__nested_attrs__',
    fallback_attrs=('inner1', 'inner2')) == [('inner3', 'z')]  # True
get_nested_objs(
    nest2,
    nested_key='does_not_exist',
    fallback_attrs='inner2') == [('inner2', 'y')]  # True

obj The object to find nested objects in.
nested_key A string that names an attribute on obj which contains the names of attributes to search for nested objects in. See function documentation for details. Default is __nested_attrs__.
fallback_attrs A string or collection of strings that name attributes to search for nested objects in, when nested_key is not present. See function documentation for details. Default is the tuple ('inner_results', 'accepted_results', 'inner_kernel'), which works well for MCMC kernels and kernel results.

pairs Returns a (possibly empty) list of (field name, nested results).