def _build_tree(value_and_gradients_fn,
                current_state,
                current_target_log_prob,
                current_grads_target_log_prob,
                current_momentum,
                direction,
                depth,
                step_size,
                log_slice_sample,
                max_simulation_error=1000.,
                seed=None):
  """Builds a tree at a given tree depth and at a given state.

  The `current` state is immediately adjacent to, but outside of,
  the subtrajectory spanned by the returned `forward` and `reverse` states.

  Args:
    value_and_gradients_fn: Python callable which takes an argument like
      `*current_state` and returns a tuple of its (possibly unnormalized)
      log-density under the target distribution and its gradient with respect to
      each state.
    current_state: List of `Tensor`s representing the current states of the
      NUTS trajectory.
    current_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `current_state`.
    current_grads_target_log_prob: List of `Tensor`s representing gradient of
      `current_target_log_prob` with respect to `current_state`. Must have same
      shape as `current_state`.
    current_momentum: List of `Tensor`s representing the momentums of
      `current_state`. Must have same shape as `current_state`.
    direction: int that is either -1 or 1. It determines whether to perform
      leapfrog integration backwards (reverse) or forward in time respectively.
    depth: non-negative int that indicates how deep of a tree to build.
      Each call to `_build_tree` takes `2**depth` leapfrog steps.
    step_size: List of `Tensor`s representing the step sizes for the leapfrog
      integrator. Must have same shape as `current_state`.
    log_slice_sample: The log of an auxiliary slice variable. It is used
      together with `max_simulation_error` to avoid simulating trajectories with
      too much numerical error.
    max_simulation_error: Maximum simulation error to tolerate before
      terminating the trajectory. Simulation error is the
      `log_slice_sample` minus the log-joint probability at the simulated state.
    seed: Integer to seed the random number generator.

  Returns:
    reverse_state: List of `Tensor`s representing the "reverse" states of the
      NUTS trajectory. Has same shape as `current_state`.
    reverse_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `reverse_state`.
    reverse_grads_target_log_prob: List of `Tensor`s representing gradient of
      `reverse_target_log_prob` with respect to `reverse_state`. Has same shape
      as `reverse_state`.
    reverse_momentum: List of `Tensor`s representing the momentums of
      `reverse_state`. Has same shape as `reverse_state`.
    forward_state: List of `Tensor`s representing the "forward" states of the
      NUTS trajectory. Has same shape as `current_state`.
    forward_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at the `forward_state`.
    forward_grads_target_log_prob: List of `Tensor`s representing gradient of
      `forward_target_log_prob` with respect to `forward_state`. Has same shape
      as `forward_state`.
    forward_momentum: List of `Tensor`s representing the momentums of
      `forward_state`. Has same shape as `forward_state`.
    next_state: List of `Tensor`s representing the next states of the NUTS
      trajectory. Has same shape as `current_state`.
    next_target_log_prob: Scalar `Tensor` representing the value of
      `target_log_prob_fn` at `next_state`.
    next_grads_target_log_prob: List of `Tensor`s representing the gradient of
      `next_target_log_prob` with respect to `next_state`.
    num_states: Number of acceptable candidate states in the subtree. A state is
      acceptable if it is "in the slice", that is, if its log-joint probability
      with its momentum is greater than `log_slice_sample`.
    continue_trajectory: bool determining whether to continue the simulation
      trajectory. The trajectory is continued if no U-turns are encountered
      within the built subtree, and if the log-probability accumulation due to
      integration error does not exceed `max_simulation_error`.
  """
  if depth == 0:  # base case
    # Take a leapfrog step. Terminate the tree-building if the simulation
    # error from the leapfrog integrator is too large. States discovered by
    # continuing the simulation are likely to have very low probability.
    [
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
    ] = _leapfrog(
        value_and_gradients_fn=value_and_gradients_fn,
        current_state=current_state,
        current_grads_target_log_prob=current_grads_target_log_prob,
        current_momentum=current_momentum,
        step_size=direction * step_size)
    next_log_joint = _log_joint(next_target_log_prob, next_momentum)
    num_states = tf.cast(next_log_joint > log_slice_sample, dtype=tf.int32)
    continue_trajectory = (next_log_joint >
                           log_slice_sample - max_simulation_error)
    return [
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        next_momentum,
        next_state,
        next_target_log_prob,
        next_grads_target_log_prob,
        num_states,
        continue_trajectory,
    ]

  # Build a tree at the current state.
  seed_stream = tfd.SeedStream(seed, "build_tree")
  [
      reverse_state,
      reverse_target_log_prob,
      reverse_grads_target_log_prob,
      reverse_momentum,
      forward_state,
      forward_target_log_prob,
      forward_grads_target_log_prob,
      forward_momentum,
      next_state,
      next_target_log_prob,
      next_grads_target_log_prob,
      num_states,
      continue_trajectory,
  ] = _build_tree(value_and_gradients_fn=value_and_gradients_fn,
                  current_state=current_state,
                  current_target_log_prob=current_target_log_prob,
                  current_grads_target_log_prob=current_grads_target_log_prob,
                  current_momentum=current_momentum,
                  direction=direction,
                  depth=depth - 1,
                  step_size=step_size,
                  log_slice_sample=log_slice_sample,
                  seed=seed_stream())
  if continue_trajectory:
    # If the just-built subtree did not terminate, build a second subtree at
    # the forward or reverse state, as appropriate.
    if direction < 0:
      [
          reverse_state,
          reverse_target_log_prob,
          reverse_grads_target_log_prob,
          reverse_momentum,
          _,
          _,
          _,
          _,
          far_state,
          far_target_log_prob,
          far_grads_target_log_prob,
          far_num_states,
          far_continue_trajectory,
      ] = _build_tree(
          value_and_gradients_fn=value_and_gradients_fn,
          current_state=reverse_state,
          current_target_log_prob=reverse_target_log_prob,
          current_grads_target_log_prob=reverse_grads_target_log_prob,
          current_momentum=reverse_momentum,
          direction=direction,
          depth=depth - 1,
          step_size=step_size,
          log_slice_sample=log_slice_sample,
          seed=seed_stream())
    else:
      [
          _,
          _,
          _,
          _,
          forward_state,
          forward_target_log_prob,
          forward_grads_target_log_prob,
          forward_momentum,
          far_state,
          far_target_log_prob,
          far_grads_target_log_prob,
          far_num_states,
          far_continue_trajectory,
      ] = _build_tree(
          value_and_gradients_fn=value_and_gradients_fn,
          current_state=forward_state,
          current_target_log_prob=forward_target_log_prob,
          current_grads_target_log_prob=forward_grads_target_log_prob,
          current_momentum=forward_momentum,
          direction=direction,
          depth=depth - 1,
          step_size=step_size,
          log_slice_sample=log_slice_sample,
          seed=seed_stream())

    # Propose either `next_state` (which came from the first subtree and so is
    # nearby) or the new forward/reverse state (which came from the second
    # subtree and so is far away).
    num_states += far_num_states
    accept_far_state = _random_bernoulli(
        [],
        probs=far_num_states / num_states,
        dtype=tf.bool,
        seed=seed_stream())
    if accept_far_state:
      next_state = far_state
      next_target_log_prob = far_target_log_prob
      next_grads_target_log_prob = far_grads_target_log_prob

    # Continue the NUTS trajectory if the far subtree did not terminate either,
    # and if the reverse-most and forward-most states do not exhibit a U-turn.
    has_no_u_turn = tf.logical_and(
        _has_no_u_turn(forward_state, reverse_state, forward_momentum),
        _has_no_u_turn(forward_state, reverse_state, reverse_momentum))
    continue_trajectory = far_continue_trajectory and has_no_u_turn

  return [
      reverse_state,
      reverse_target_log_prob,
      reverse_grads_target_log_prob,
      reverse_momentum,
      forward_state,
      forward_target_log_prob,
      forward_grads_target_log_prob,
      forward_momentum,
      next_state,
      next_target_log_prob,
      next_grads_target_log_prob,
      num_states,
      continue_trajectory,
  ]

Builds a tree at a given tree depth and at a given state .


def _embed_no_none_gradient_check(value_and_gradients_fn):
  """Wraps value and gradients function to assist with None gradients."""
  @functools.wraps(value_and_gradients_fn)
  def func_wrapped(*args, **kwargs):
    """Wrapped function which checks for None gradients."""
    value, grads = value_and_gradients_fn(*args, **kwargs)
    if any(grad is None for grad in grads):
      raise ValueError("Gradient is None for a state.")
    return value, grads
  return func_wrapped

def
_embed_no_none_gradient_check
(
value_and_gradients_fn
)
:
@
functools
.
wraps
(
value_and_gradients_fn
)
def
func_wrapped
(
*
args
,
**
kwargs
)
:
value
,
grads
=
value_and_gradients_fn
(
*
args
,
**
kwargs
)
if
any
(
grad
is
for
grad
in
grads
)
:
raise
ValueError
(
"Gradient is None for a state."
)
return
value
,
grads
return
func_wrapped

Wraps value and gradients function to assist with None gradients .


def _has_no_u_turn(state_one, state_two, momentum):
  """If two given states and momentum do not exhibit a U-turn pattern."""
  dot_product = sum([
      tf.reduce_sum(input_tensor=(s1 - s2) * m)
      for s1, s2, m in zip(state_one, state_two, momentum)
  ])
  return dot_product > 0

def _has_no_u_turn ( state_one , state_two , momentum ) : dot_product = sum ( [ tf . reduce_sum ( input_tensor = ( s1 - s2 ) * m ) for s1 , s2 , m in zip ( state_one , state_two , momentum ) ] ) return dot_product > 0

If two given states and momentum do not exhibit a U - turn pattern .


def _leapfrog(value_and_gradients_fn,
              current_state,
              current_grads_target_log_prob,
              current_momentum,
              step_size):
  """Runs one step of leapfrog integration."""
  mid_momentum = [
      m + 0.5 * step * g for m, step, g in
      zip(current_momentum, step_size, current_grads_target_log_prob)]
  next_state = [
      s + step * m for s, step, m in
      zip(current_state, step_size, mid_momentum)]
  next_target_log_prob, next_grads_target_log_prob = value_and_gradients_fn(
      *next_state)
  next_momentum = [
      m + 0.5 * step * g for m, step, g in
      zip(mid_momentum, step_size, next_grads_target_log_prob)]
  return [
      next_state,
      next_target_log_prob,
      next_grads_target_log_prob,
      next_momentum,
  ]

def _leapfrog ( value_and_gradients_fn , current_state , current_grads_target_log_prob , current_momentum , step_size ) : mid_momentum = [ m + 0.5 * step * g for m , step , g in zip ( current_momentum , step_size , current_grads_target_log_prob ) ] next_state = [ s + step * m for s , step , m in zip ( current_state , step_size , mid_momentum ) ] next_target_log_prob , next_grads_target_log_prob = value_and_gradients_fn ( * next_state ) next_momentum = [ m + 0.5 * step * g for m , step , g in zip ( mid_momentum , step_size , next_grads_target_log_prob ) ] return [ next_state , next_target_log_prob , next_grads_target_log_prob , next_momentum , ]


def _log_joint(current_target_log_prob, current_momentum):
  """Log-joint probability given a state's log-probability and momentum."""
  momentum_log_prob = -sum(
      [tf.reduce_sum(input_tensor=0.5 * (m**2.)) for m in current_momentum])
  return current_target_log_prob + momentum_log_prob

def _log_joint ( current_target_log_prob , current_momentum ) : momentum_log_prob = - sum ( [ tf . reduce_sum ( input_tensor = 0.5 * ( m ** 2. ) ) for m in current_momentum ] ) return current_target_log_prob + momentum_log_prob

Log - joint probability given a state s log - probability and momentum .


def _random_bernoulli(shape, probs, dtype=tf.int32, seed=None, name=None):
  """Returns samples from a Bernoulli distribution."""
  with tf.compat.v1.name_scope(name, "random_bernoulli", [shape, probs]):
    probs = tf.convert_to_tensor(value=probs)
    random_uniform = tf.random.uniform(shape, dtype=probs.dtype, seed=seed)
    return tf.cast(tf.less(random_uniform, probs), dtype)

def _random_bernoulli ( shape , probs , dtype = tf . int32 , seed = None , name = None ) : with tf . compat . v1 . name_scope ( name , "random_bernoulli" , [ shape , probs ] ) : probs = tf . convert_to_tensor ( value = probs ) random_uniform = tf . random . uniform ( shape , dtype = probs . dtype , seed = seed ) return tf . cast ( tf . less ( random_uniform , probs ) , dtype )

Returns samples from a Bernoulli distribution .


def default_loc_scale_fn(
    is_singular=False,
    loc_initializer=tf.compat.v1.initializers.random_normal(stddev=0.1),
    untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(
        mean=-3., stddev=0.1),
    loc_regularizer=None,
    untransformed_scale_regularizer=None,
    loc_constraint=None,
    untransformed_scale_constraint=None):
  """Makes closure which creates `loc`, `scale` params from `tf.get_variable`.

  This function produces a closure which produces `loc`, `scale` using
  `tf.get_variable`. The closure accepts the following arguments:

    dtype: Type of parameter's event.
    shape: Python `list`-like representing the parameter's event shape.
    name: Python `str` name prepended to any created (or existing)
      `tf.Variable`s.
    trainable: Python `bool` indicating all created `tf.Variable`s should be
      added to the graph collection `GraphKeys.TRAINABLE_VARIABLES`.
    add_variable_fn: `tf.get_variable`-like `callable` used to create (or
      access existing) `tf.Variable`s.

  Args:
    is_singular: Python `bool` indicating if `scale is None`. Default: `False`.
    loc_initializer: Initializer function for the `loc` parameters.
      The default is `tf.random_normal_initializer(mean=0., stddev=0.1)`.
    untransformed_scale_initializer: Initializer function for the `scale`
      parameters. Default value: `tf.random_normal_initializer(mean=-3.,
      stddev=0.1)`. This implies the softplus transformed result is initialized
      near `0`. It allows a `Normal` distribution with `scale` parameter set to
      this value to approximately act like a point mass.
    loc_regularizer: Regularizer function for the `loc` parameters.
      The default (`None`) is to use the `tf.get_variable` default.
    untransformed_scale_regularizer: Regularizer function for the `scale`
      parameters. The default (`None`) is to use the `tf.get_variable` default.
    loc_constraint: An optional projection function to be applied to the
      loc after being updated by an `Optimizer`. The function must take as input
      the unprojected variable and must return the projected variable (which
      must have the same shape). Constraints are not safe to use when doing
      asynchronous distributed training.
      The default (`None`) is to use the `tf.get_variable` default.
    untransformed_scale_constraint: An optional projection function to be
      applied to the `scale` parameters after being updated by an `Optimizer`
      (e.g. used to implement norm constraints or value constraints). The
      function must take as input the unprojected variable and must return the
      projected variable (which must have the same shape). Constraints are not
      safe to use when doing asynchronous distributed training. The default
      (`None`) is to use the `tf.get_variable` default.

  Returns:
    default_loc_scale_fn: Python `callable` which instantiates `loc`, `scale`
    parameters from args: `dtype, shape, name, trainable, add_variable_fn`.
  """
  def _fn(dtype, shape, name, trainable, add_variable_fn):
    """Creates `loc`, `scale` parameters."""
    loc = add_variable_fn(
        name=name + '_loc',
        shape=shape,
        initializer=loc_initializer,
        regularizer=loc_regularizer,
        constraint=loc_constraint,
        dtype=dtype,
        trainable=trainable)
    if is_singular:
      return loc, None
    untransformed_scale = add_variable_fn(
        name=name + '_untransformed_scale',
        shape=shape,
        initializer=untransformed_scale_initializer,
        regularizer=untransformed_scale_regularizer,
        constraint=untransformed_scale_constraint,
        dtype=dtype,
        trainable=trainable)
    scale = (np.finfo(dtype.as_numpy_dtype).eps +
             tf.nn.softplus(untransformed_scale))
    return loc, scale
  return _fn

def default_loc_scale_fn ( is_singular = False , loc_initializer = tf . compat . v1 . initializers . random_normal ( stddev = 0.1 ) , untransformed_scale_initializer = tf . compat . v1 . initializers . random_normal ( mean = - 3. , stddev = 0.1 ) , loc_regularizer = None , untransformed_scale_regularizer = None , loc_constraint = None , untransformed_scale_constraint = None ) : def _fn ( dtype , shape , name , trainable , add_variable_fn ) : loc = add_variable_fn ( name = name + '_loc' , shape = shape , initializer = loc_initializer , regularizer = loc_regularizer , constraint = loc_constraint , dtype = dtype , trainable = trainable ) if is_singular : return loc , None untransformed_scale = add_variable_fn ( name = name + '_untransformed_scale' , shape = shape , initializer = untransformed_scale_initializer , regularizer = untransformed_scale_regularizer , constraint = untransformed_scale_constraint , dtype = dtype , trainable = trainable ) scale = ( np . finfo ( dtype . as_numpy_dtype ) . eps + tf . nn . softplus ( untransformed_scale ) ) return loc , scale return _fn

Makes closure which creates loc scale params from tf . get_variable .


def default_mean_field_normal_fn(
    is_singular=False,
    loc_initializer=tf.compat.v1.initializers.random_normal(stddev=0.1),
    untransformed_scale_initializer=tf.compat.v1.initializers.random_normal(
        mean=-3., stddev=0.1),
    loc_regularizer=None,
    untransformed_scale_regularizer=None,
    loc_constraint=None,
    untransformed_scale_constraint=None):
  """Creates a function to build Normal distributions with trainable params.

  This function produces a closure which produces `tfd.Normal`
  parameterized by a loc` and `scale` each created using `tf.get_variable`.

  Args:
    is_singular: Python `bool` if `True`, forces the special case limit of
      `scale->0`, i.e., a `Deterministic` distribution.
    loc_initializer: Initializer function for the `loc` parameters.
      The default is `tf.random_normal_initializer(mean=0., stddev=0.1)`.
    untransformed_scale_initializer: Initializer function for the `scale`
      parameters. Default value: `tf.random_normal_initializer(mean=-3.,
      stddev=0.1)`. This implies the softplus transformed result is initialized
      near `0`. It allows a `Normal` distribution with `scale` parameter set to
      this value to approximately act like a point mass.
    loc_regularizer: Regularizer function for the `loc` parameters.
    untransformed_scale_regularizer: Regularizer function for the `scale`
      parameters.
    loc_constraint: An optional projection function to be applied to the
      loc after being updated by an `Optimizer`. The function must take as input
      the unprojected variable and must return the projected variable (which
      must have the same shape). Constraints are not safe to use when doing
      asynchronous distributed training.
    untransformed_scale_constraint: An optional projection function to be
      applied to the `scale` parameters after being updated by an `Optimizer`
      (e.g. used to implement norm constraints or value constraints). The
      function must take as input the unprojected variable and must return the
      projected variable (which must have the same shape). Constraints are not
      safe to use when doing asynchronous distributed training.

  Returns:
    make_normal_fn: Python `callable` which creates a `tfd.Normal`
      using from args: `dtype, shape, name, trainable, add_variable_fn`.
  """
  loc_scale_fn = default_loc_scale_fn(
      is_singular=is_singular,
      loc_initializer=loc_initializer,
      untransformed_scale_initializer=untransformed_scale_initializer,
      loc_regularizer=loc_regularizer,
      untransformed_scale_regularizer=untransformed_scale_regularizer,
      loc_constraint=loc_constraint,
      untransformed_scale_constraint=untransformed_scale_constraint)
  def _fn(dtype, shape, name, trainable, add_variable_fn):
    """Creates multivariate `Deterministic` or `Normal` distribution.

    Args:
      dtype: Type of parameter's event.
      shape: Python `list`-like representing the parameter's event shape.
      name: Python `str` name prepended to any created (or existing)
        `tf.Variable`s.
      trainable: Python `bool` indicating all created `tf.Variable`s should be
        added to the graph collection `GraphKeys.TRAINABLE_VARIABLES`.
      add_variable_fn: `tf.get_variable`-like `callable` used to create (or
        access existing) `tf.Variable`s.

    Returns:
      Multivariate `Deterministic` or `Normal` distribution.
    """
    loc, scale = loc_scale_fn(dtype, shape, name, trainable, add_variable_fn)
    if scale is None:
      dist = tfd.Deterministic(loc=loc)
    else:
      dist = tfd.Normal(loc=loc, scale=scale)
    batch_ndims = tf.size(input=dist.batch_shape_tensor())
    return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)
  return _fn

def default_mean_field_normal_fn ( is_singular = False , loc_initializer = tf . compat . v1 . initializers . random_normal ( stddev = 0.1 ) , untransformed_scale_initializer = tf . compat . v1 . initializers . random_normal ( mean = - 3. , stddev = 0.1 ) , loc_regularizer = None , untransformed_scale_regularizer = None , loc_constraint = None , untransformed_scale_constraint = None ) : loc_scale_fn = default_loc_scale_fn ( is_singular = is_singular , loc_initializer = loc_initializer , untransformed_scale_initializer = untransformed_scale_initializer , loc_regularizer = loc_regularizer , untransformed_scale_regularizer = untransformed_scale_regularizer , loc_constraint = loc_constraint , untransformed_scale_constraint = untransformed_scale_constraint ) def _fn ( dtype , shape , name , trainable , add_variable_fn ) : loc , scale = loc_scale_fn ( dtype , shape , name , trainable , add_variable_fn ) if scale is None : dist = tfd . Deterministic ( loc = loc ) else : dist = tfd . Normal ( loc = loc , scale = scale ) batch_ndims = tf . size ( input = dist . batch_shape_tensor ( ) ) return tfd . Independent ( dist , reinterpreted_batch_ndims = batch_ndims ) return _fn

Creates a function to build Normal distributions with trainable params .


def default_multivariate_normal_fn(dtype, shape, name, trainable,
                                   add_variable_fn):
  """Creates multivariate standard `Normal` distribution.

  Args:
    dtype: Type of parameter's event.
    shape: Python `list`-like representing the parameter's event shape.
    name: Python `str` name prepended to any created (or existing)
      `tf.Variable`s.
    trainable: Python `bool` indicating all created `tf.Variable`s should be
      added to the graph collection `GraphKeys.TRAINABLE_VARIABLES`.
    add_variable_fn: `tf.get_variable`-like `callable` used to create (or
      access existing) `tf.Variable`s.

  Returns:
    Multivariate standard `Normal` distribution.
  """
  del name, trainable, add_variable_fn   # unused
  dist = tfd.Normal(loc=tf.zeros(shape, dtype), scale=dtype.as_numpy_dtype(1))
  batch_ndims = tf.size(input=dist.batch_shape_tensor())
  return tfd.Independent(dist, reinterpreted_batch_ndims=batch_ndims)

def default_multivariate_normal_fn ( dtype , shape , name , trainable , add_variable_fn ) : del name , trainable , add_variable_fn dist = tfd . Normal ( loc = tf . zeros ( shape , dtype ) , scale = dtype . as_numpy_dtype ( 1 ) ) batch_ndims = tf . size ( input = dist . batch_shape_tensor ( ) ) return tfd . Independent ( dist , reinterpreted_batch_ndims = batch_ndims )

Creates multivariate standard Normal distribution .


def deserialize_function(serial, function_type):
  """Deserializes the Keras-serialized function.

  (De)serializing Python functions from/to bytecode is unsafe. Therefore we
  also use the function's type as an anonymous function ('lambda') or named
  function in the Python environment ('function'). In the latter case, this lets
  us use the Python scope to obtain the function rather than reload it from
  bytecode. (Note that both cases are brittle!)

  Keras-deserialized functions do not perform lexical scoping. Any modules that
  the function requires must be imported within the function itself.

  This serialization mimicks the implementation in `tf.keras.layers.Lambda`.

  Args:
    serial: Serialized Keras object: typically a dict, string, or bytecode.
    function_type: Python string denoting 'function' or 'lambda'.

  Returns:
    function: Function the serialized Keras object represents.

  #### Examples

  ```python
  serial, function_type = serialize_function(lambda x: x)
  function = deserialize_function(serial, function_type)
  assert function(2.3) == 2.3  # function is identity
  ```

  """
  if function_type == 'function':
    # Simple lookup in custom objects
    function = tf.keras.utils.deserialize_keras_object(serial)
  elif function_type == 'lambda':
    # Unsafe deserialization from bytecode
    function = generic_utils.func_load(serial)
  else:
    raise TypeError('Unknown function type:', function_type)
  return function

def
deserialize_function
(
serial
,
function_type
)
:
if
function_type
==
'function'
:
function
=
tf
.
keras
.
utils
.
deserialize_keras_object
(
serial
)
elif
function_type
==
'lambda'
:
function
=
generic_utils
.
func_load
(
serial
)
else
:
raise
TypeError
(
'Unknown function type:'
,
function_type
)
return
function

Deserializes the Keras - serialized function .


def serialize_function(func):
  """Serializes function for Keras.

  (De)serializing Python functions from/to bytecode is unsafe. Therefore we
  return the function's type as an anonymous function ('lambda') or named
  function in the Python environment ('function'). In the latter case, this lets
  us use the Python scope to obtain the function rather than reload it from
  bytecode. (Note that both cases are brittle!)

  This serialization mimicks the implementation in `tf.keras.layers.Lambda`.

  Args:
    func: Python function to serialize.

  Returns:
    (serial, function_type): Serialized object, which is a tuple of its
    bytecode (if function is anonymous) or name (if function is named), and its
    function type.
  """
  if isinstance(func, types.LambdaType):
    return generic_utils.func_dump(func), 'lambda'
  return func.__name__, 'function'

def serialize_function ( func ) : if isinstance ( func , types . LambdaType ) : return generic_utils . func_dump ( func ) , 'lambda' return func . __name__ , 'function'


def broadcast_structure(to_structure, from_structure):
  """Broadcasts `from_structure` to `to_structure`.

  This is useful for downstream usage of `zip` or `tf.nest.map_structure`.

  If `from_structure` is a singleton, it is tiled to match the structure of
  `to_structure`. Note that the elements in `from_structure` are not copied if
  this tiling occurs.

  Args:
    to_structure: A structure.
    from_structure: A structure.

  Returns:
    new_from_structure: Same structure as `to_structure`.

  #### Example:

  ```python
  a_structure = ['a', 'b', 'c']
  b_structure = broadcast_structure(a_structure, 'd')
  # -> ['d', 'd', 'd']
  c_structure = tf.nest.map_structure(
      lambda a, b: a + b, a_structure, b_structure)
  # -> ['ad', 'bd', 'cd']
  ```
  """
  from_parts = tf.nest.flatten(from_structure)
  if len(from_parts) == 1:
    from_structure = tf.nest.map_structure(lambda _: from_parts[0],
                                           to_structure)
  return from_structure

def broadcast_structure ( to_structure , from_structure ) : from_parts = tf . nest . flatten ( from_structure ) if len ( from_parts ) == 1 : from_structure = tf . nest . map_structure ( lambda _ : from_parts [ 0 ] , to_structure ) return from_structure


def _nested_convert_to_tensor(struct, dtype=None, name=None):
  """Eagerly converts struct to Tensor, recursing upon failure."""
  if dtype is not None or not tf.nest.is_nested(struct):
    return tf.convert_to_tensor(struct, dtype=dtype)

  if _maybe_convertible_to_tensor(struct):
    try:
      # Try converting the structure wholesale.
      return tf.convert_to_tensor(value=struct, name=name)
    except (ValueError, TypeError):
      # Unfortunately Eager/Graph mode don't agree on the error type.
      pass
  # Try converting all of its children.
  shallow_struct = _get_shallow_structure(struct)
  return nest.map_structure_up_to(
      shallow_struct, lambda s: _nested_convert_to_tensor(s, name=name), struct)

def _nested_convert_to_tensor ( struct , dtype = None , name = None ) : if dtype is not None or not tf . nest . is_nested ( struct ) : return tf . convert_to_tensor ( struct , dtype = dtype ) if _maybe_convertible_to_tensor ( struct ) : try : return tf . convert_to_tensor ( value = struct , name = name ) except ( ValueError , TypeError ) : pass shallow_struct = _get_shallow_structure ( struct ) return nest . map_structure_up_to ( shallow_struct , lambda s : _nested_convert_to_tensor ( s , name = name ) , struct )

Eagerly converts struct to Tensor recursing upon failure .


def convert_args_to_tensor(args, dtype=None, name=None):
  """Converts `args` to `Tensor`s.

  Use this when it is necessary to convert user-provided arguments that will
  then be passed to user-provided callables.

  When `dtype` is `None` this function behaves as follows:

  1A. If the top-level structure is a `list`/`tuple` but not a `namedtuple`,
      then it is left as is and only its elements are converted to `Tensor`s.

  2A. The sub-structures are converted to `Tensor`s eagerly. E.g. if `args` is
      `{'arg': [[1], [2]]}` it is converted to
      `{'arg': tf.constant([[1], [2]])}`. If the conversion fails, it will
      attempt to recurse into its children.

  When `dtype` is specified, it acts as both a structural and numeric type
  constraint. `dtype` can be a single `DType`, `None` or a nested collection
  thereof. The conversion rule becomes as follows:

  1B. The return value of this function will have the same structure as `dtype`.

  2B. If the leaf of `dtype` is a concrete `DType`, then the corresponding
      sub-structure in `args` is converted to a `Tensor`.

  3B. If the leaf of `dtype` is `None`, then the corresponding sub-structure is
      converted eagerly as described in the rule 2A above.

  Args:
    args: Arguments to convert to `Tensor`s.
    dtype: Optional structure/numeric type constraint.
    name: Optional name-scope to use.

  Returns:
    args: Converted `args`.

  #### Examples.

  This table shows some useful conversion cases. `T` means `Tensor`, `NT` means
  `namedtuple` and `CNT` means a `namedtuple` with a `Tensor`-conversion
  function registered.

  |     args     |    dtype   |       output       |
  |:------------:|:----------:|:------------------:|
  | `{"a": 1}`   | `None`     | `{"a": T(1)}`      |
  | `T(1)`       | `None`     | `T(1)`             |
  | `[1]`        | `None`     | `[T(1)]`           |
  | `[1]`        | `tf.int32` | `T([1])`           |
  | `[[T(1)]]`   | `None`     | `[T([1])]`         |
  | `[[T(1)]]`   | `[[None]]` | `[[T(1)]]`         |
  | `NT(1, 2)`   | `None`     | `NT(T(1), T(2))`   |
  | `NT(1, 2)`   | `tf.int32` | `T([1, 2])`        |
  | `CNT(1, 2)`  | `None`     | `T(...)`           |
  | `[[1, [2]]]` | `None`     | `[[T(1), T([2])]]` |

  """
  if dtype is None:
    if expand_as_args(args) or _expand_as_kwargs(args):
      shallow_args = _get_shallow_structure(args)
      return nest.map_structure_up_to(
          shallow_args, lambda s: _nested_convert_to_tensor(s, name=name), args)
    else:
      return _nested_convert_to_tensor(args, name=name)
  else:
    return nest.map_structure_up_to(
        dtype, lambda s, dtype: _nested_convert_to_tensor(s, dtype, name), args,
        dtype)

def convert_args_to_tensor ( args , dtype = None , name = None ) : if dtype is None : if expand_as_args ( args ) or _expand_as_kwargs ( args ) : shallow_args = _get_shallow_structure ( args ) return nest . map_structure_up_to ( shallow_args , lambda s : _nested_convert_to_tensor ( s , name = name ) , args ) else : return _nested_convert_to_tensor ( args , name = name ) else : return nest . map_structure_up_to ( dtype , lambda s , dtype : _nested_convert_to_tensor ( s , dtype , name ) , args , dtype )


def call_fn(fn, args):
  """Calls `fn` with `args`, possibly expanding `args`.

  Use this function when calling a user-provided callable using user-provided
  arguments.

  The expansion rules are as follows:

  `fn(*args)` if `args` is a `list` or a `tuple`, but not a `namedtuple`.
  `fn(**args)` if `args` is a `dict`.
  `fn(args)` otherwise.

  Args:
    fn: A callable that takes either `args` as an argument(s).
    args: Arguments to `fn`.

  Returns:
    result: Return value of `fn`.
  """

  if expand_as_args(args):
    return fn(*args)
  elif _expand_as_kwargs(args):
    return fn(**args)
  else:
    return fn(args)

def call_fn ( fn , args ) : if expand_as_args ( args ) : return fn ( * args ) elif _expand_as_kwargs ( args ) : return fn ( ** args ) else : return fn ( args )

Calls fn with args possibly expanding args .


def _get_tensor_like_attributes():
  """Returns `Tensor` attributes related to shape and Python builtins."""
  # Enable "Tensor semantics" for distributions.
  # See tensorflow/python/framework/ops.py `class Tensor` for details.
  attrs = dict()
  # Setup overloadable operators and white-listed members / properties.
  attrs.update((attr, _wrap_method(tf.Tensor, attr))
               for attr in tf.Tensor.OVERLOADABLE_OPERATORS.union({'__iter__'}))
  # Copy some members straight-through.
  attrs.update((attr, getattr(tf.Tensor, attr))
               for attr in {'__nonzero__', '__bool__', '__array_priority__'})
  return attrs

def _get_tensor_like_attributes ( ) : attrs = dict ( ) attrs . update ( ( attr , _wrap_method ( tf . Tensor , attr ) ) for attr in tf . Tensor . OVERLOADABLE_OPERATORS . union ( { '__iter__' } ) ) attrs . update ( ( attr , getattr ( tf . Tensor , attr ) ) for attr in { '__nonzero__' , '__bool__' , '__array_priority__' } ) return attrs

Returns Tensor attributes related to shape and Python builtins .


def make_mixture_prior(latent_size, mixture_components):
  """Creates the mixture of Gaussians prior distribution.

  Args:
    latent_size: The dimensionality of the latent representation.
    mixture_components: Number of elements of the mixture.

  Returns:
    random_prior: A `tfd.Distribution` instance representing the distribution
      over encodings in the absence of any evidence.
  """
  if mixture_components == 1:
    # See the module docstring for why we don't learn the parameters here.
    return tfd.MultivariateNormalDiag(
        loc=tf.zeros([latent_size]),
        scale_identity_multiplier=1.0)

  loc = tf.compat.v1.get_variable(
      name="loc", shape=[mixture_components, latent_size])
  raw_scale_diag = tf.compat.v1.get_variable(
      name="raw_scale_diag", shape=[mixture_components, latent_size])
  mixture_logits = tf.compat.v1.get_variable(
      name="mixture_logits", shape=[mixture_components])

  return tfd.MixtureSameFamily(
      components_distribution=tfd.MultivariateNormalDiag(
          loc=loc,
          scale_diag=tf.nn.softplus(raw_scale_diag)),
      mixture_distribution=tfd.Categorical(logits=mixture_logits),
      name="prior")

def make_mixture_prior ( latent_size , mixture_components ) : if mixture_components == 1 : return tfd . MultivariateNormalDiag ( loc = tf . zeros ( [ latent_size ] ) , scale_identity_multiplier = 1.0 ) loc = tf . compat . v1 . get_variable ( name = "loc" , shape = [ mixture_components , latent_size ] ) raw_scale_diag = tf . compat . v1 . get_variable ( name = "raw_scale_diag" , shape = [ mixture_components , latent_size ] ) mixture_logits = tf . compat . v1 . get_variable ( name = "mixture_logits" , shape = [ mixture_components ] ) return tfd . MixtureSameFamily ( components_distribution = tfd . MultivariateNormalDiag ( loc = loc , scale_diag = tf . nn . softplus ( raw_scale_diag ) ) , mixture_distribution = tfd . Categorical ( logits = mixture_logits ) , name = "prior" )

Creates the mixture of Gaussians prior distribution .


def pack_images(images, rows, cols):
  """Helper utility to make a field of images."""
  shape = tf.shape(input=images)
  width = shape[-3]
  height = shape[-2]
  depth = shape[-1]
  images = tf.reshape(images, (-1, width, height, depth))
  batch = tf.shape(input=images)[0]
  rows = tf.minimum(rows, batch)
  cols = tf.minimum(batch // rows, cols)
  images = images[:rows * cols]
  images = tf.reshape(images, (rows, cols, width, height, depth))
  images = tf.transpose(a=images, perm=[0, 2, 1, 3, 4])
  images = tf.reshape(images, [1, rows * width, cols * height, depth])
  return images

def pack_images ( images , rows , cols ) : shape = tf . shape ( input = images ) width = shape [ - 3 ] height = shape [ - 2 ] depth = shape [ - 1 ] images = tf . reshape ( images , ( - 1 , width , height , depth ) ) batch = tf . shape ( input = images ) [ 0 ] rows = tf . minimum ( rows , batch ) cols = tf . minimum ( batch // rows , cols ) images = images [ : rows * cols ] images = tf . reshape ( images , ( rows , cols , width , height , depth ) ) images = tf . transpose ( a = images , perm = [ 0 , 2 , 1 , 3 , 4 ] ) images = tf . reshape ( images , [ 1 , rows * width , cols * height , depth ] ) return images

Helper utility to make a field of images .


def download(directory, filename):
  """Downloads a file."""
  filepath = os.path.join(directory, filename)
  if tf.io.gfile.exists(filepath):
    return filepath
  if not tf.io.gfile.exists(directory):
    tf.io.gfile.makedirs(directory)
  url = os.path.join(ROOT_PATH, filename)
  print("Downloading %s to %s" % (url, filepath))
  urllib.request.urlretrieve(url, filepath)
  return filepath

def
download
(
directory
,
filename
)
:
filepath
=
os
.
path
.
join
(
directory
,
filename
)
if
tf
.
io
.
gfile
.
exists
(
filepath
)
:
return
filepath
if
not
tf
.
io
.
gfile
.
exists
(
directory
)
:
tf
.
io
.
gfile
.
makedirs
(
directory
)
url
=
os
.
path
.
join
(
ROOT_PATH
,
filename
)
print
(
"Downloading %s to %s"
%
(
url
,
filepath
)
)
urllib
.
request
.
urlretrieve
(
url
,
filepath
)
return
filepath


def build_fake_input_fns(batch_size):
  """Builds fake MNIST-style data for unit testing."""
  random_sample = np.random.rand(batch_size, *IMAGE_SHAPE).astype("float32")

  def train_input_fn():
    dataset = tf.data.Dataset.from_tensor_slices(
        random_sample).map(lambda row: (row, 0)).batch(batch_size).repeat()
    return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

  def eval_input_fn():
    dataset = tf.data.Dataset.from_tensor_slices(
        random_sample).map(lambda row: (row, 0)).batch(batch_size)
    return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

  return train_input_fn, eval_input_fn

def build_fake_input_fns ( batch_size ) : random_sample = np . random . rand ( batch_size , * IMAGE_SHAPE ) . astype ( "float32" ) def train_input_fn ( ) : dataset = tf . data . Dataset . from_tensor_slices ( random_sample ) . map ( lambda row : ( row , 0 ) ) . batch ( batch_size ) . repeat ( ) return tf . compat . v1 . data . make_one_shot_iterator ( dataset ) . get_next ( ) def eval_input_fn ( ) : dataset = tf . data . Dataset . from_tensor_slices ( random_sample ) . map ( lambda row : ( row , 0 ) ) . batch ( batch_size ) return tf . compat . v1 . data . make_one_shot_iterator ( dataset ) . get_next ( ) return train_input_fn , eval_input_fn

Builds fake MNIST - style data for unit testing .


def _validate_block_sizes(block_sizes, bijectors, validate_args):
  """Helper to validate block sizes."""
  block_sizes_shape = block_sizes.shape
  if tensorshape_util.is_fully_defined(block_sizes_shape):
    if (tensorshape_util.rank(block_sizes_shape) != 1 or
        (tensorshape_util.num_elements(block_sizes_shape) != len(bijectors))):
      raise ValueError(
          '`block_sizes` must be `None`, or a vector of the same length as '
          '`bijectors`. Got a `Tensor` with shape {} and `bijectors` of '
          'length {}'.format(block_sizes_shape, len(bijectors)))
    return block_sizes
  elif validate_args:
    message = ('`block_sizes` must be `None`, or a vector of the same length '
               'as `bijectors`.')
    with tf.control_dependencies([
        assert_util.assert_equal(
            tf.size(input=block_sizes), len(bijectors), message=message),
        assert_util.assert_equal(tf.rank(block_sizes), 1)
    ]):
      return tf.identity(block_sizes)
  else:
    return block_sizes


def maybe_check_wont_broadcast(flat_xs, validate_args):
  """Verifies that `parts` don't broadcast."""
  flat_xs = tuple(flat_xs)  # So we can receive generators.
  if not validate_args:
    # Note: we don't try static validation because it is theoretically
    # possible that a user wants to take advantage of broadcasting.
    # Only when `validate_args` is `True` do we enforce the validation.
    return flat_xs
  msg = 'Broadcasting probably indicates an error in model specification.'
  s = tuple(x.shape for x in flat_xs)
  if all(tensorshape_util.is_fully_defined(s_) for s_ in s):
    if not all(a == b for a, b in zip(s[1:], s[:-1])):
      raise ValueError(msg)
    return flat_xs
  assertions = [assert_util.assert_equal(a, b, message=msg)
                for a, b in zip(s[1:], s[:-1])]
  with tf.control_dependencies(assertions):
    return tuple(tf.identity(x) for x in flat_xs)


def multivariate_normal_tril(x,
                             dims,
                             layer_fn=tf.compat.v1.layers.dense,
                             loc_fn=lambda x: x,
                             scale_fn=tril_with_diag_softplus_and_shift,
                             name=None):
  """Constructs a trainable `tfd.MultivariateNormalTriL` distribution.

  This function creates a MultivariateNormal (MVN) with lower-triangular scale
  matrix. By default the MVN is parameterized via affine transformation of input
  tensor `x`. Using default args, this function is mathematically equivalent to:

  ```none
  Y = MVN(loc=matmul(W, x) + b,
          scale_tril=f(reshape_tril(matmul(M, x) + c)))

  where,
    W in R^[d, n]
    M in R^[d*(d+1)/2, n]
    b in R^d
    c in R^d
    f(S) = set_diag(S, softplus(matrix_diag_part(S)) + 1e-5)
  ```

  Observe that `f` makes the diagonal of the triangular-lower scale matrix be
  positive and no smaller than `1e-5`.

  #### Examples

  ```python
  # This example fits a multilinear regression loss.
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  y_size = 2  # size of single y
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size, y_size).astype(dtype)
    b = np.random.randn(1, y_size).astype(dtype)
    true_mean = np.tensordot(x, w, axes=[[-1], [0]]) + b
    noise = np.random.randn(n, y_size).astype(dtype)
    y = true_mean + noise
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting MVNTriL maximum likelihood estimator.
  mvn = tfp.trainable_distributions.multivariate_normal_tril(x, dims=y_size)
  loss = -tf.reduce_mean(mvn.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-3).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, mvn.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:38.2020797729  mse:4.17175960541
  #     iteration:200  loss:2.90179634094  mse:0.990987896919
  #     iteration:400  loss:2.82727336884  mse:0.990926623344
  #     iteration:600  loss:2.82726788521  mse:0.990926682949
  #     iteration:800  loss:2.82726788521  mse:0.990926682949
  #     iteration:999  loss:2.82726788521  mse:0.990926682949
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    dims: Scalar, `int`, `Tensor` indicated the MVN event size, i.e., the
      created MVN will be distribution over length-`dims` vectors.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [d]], axis=0)`.
      Default value: `tf.layers.dense`.
    loc_fn: Python `callable` which transforms the `loc` parameter. Takes a
      (batch of) length-`dims` vectors and returns a `Tensor` of same shape and
      `dtype`.
      Default value: `lambda x: x`.
    scale_fn: Python `callable` which transforms the `scale` parameters. Takes a
      (batch of) length-`dims * (dims + 1) / 2` vectors and returns a
      lower-triangular `Tensor` of same batch shape with rightmost dimensions
      having shape `[dims, dims]`.
      Default value: `tril_with_diag_softplus_and_shift`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "multivariate_normal_tril").

  Returns:
    mvntril: An instance of `tfd.MultivariateNormalTriL`.
  """
  with tf.compat.v1.name_scope(name, 'multivariate_normal_tril', [x, dims]):
    x = tf.convert_to_tensor(value=x, name='x')
    x = layer_fn(x, dims + dims * (dims + 1) // 2)
    return tfd.MultivariateNormalTriL(
        loc=loc_fn(x[..., :dims]),
        scale_tril=scale_fn(x[..., dims:]))

def multivariate_normal_tril ( x , dims , layer_fn = tf . compat . v1 . layers . dense , loc_fn = lambda x : x , scale_fn = tril_with_diag_softplus_and_shift , name = None ) : with tf . compat . v1 . name_scope ( name , 'multivariate_normal_tril' , [ x , dims ] ) : x = tf . convert_to_tensor ( value = x , name = 'x' ) x = layer_fn ( x , dims + dims * ( dims + 1 ) // 2 ) return tfd . MultivariateNormalTriL ( loc = loc_fn ( x [ ... , : dims ] ) , scale_tril = scale_fn ( x [ ... , dims : ] ) )

Constructs a trainable tfd . MultivariateNormalTriL distribution .


def bernoulli(x, layer_fn=tf.compat.v1.layers.dense, name=None):
  """Constructs a trainable `tfd.Bernoulli` distribution.

  This function creates a Bernoulli distribution parameterized by logits.
  Using default args, this function is mathematically equivalent to:

  ```none
  Y = Bernoulli(logits=matmul(W, x) + b)

  where,
    W in R^[d, n]
    b in R^d
  ```

  #### Examples

  This function can be used as a [logistic regression](
  https://en.wikipedia.org/wiki/Logistic_regression) loss.

  ```python
  # This example fits a logistic regression loss.
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size).astype(dtype)
    b = np.random.randn(1).astype(dtype)
    true_logits = np.tensordot(x, w, axes=[[-1], [-1]]) + b
    noise = np.random.logistic(size=n).astype(dtype)
    y = dtype(true_logits + noise > 0.)
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting Bernoulli maximum likelihood estimator.
  bernoulli = tfp.trainable_distributions.bernoulli(x)
  loss = -tf.reduce_mean(bernoulli.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, bernoulli.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:0.635675370693  mse:0.222526371479
  #     iteration:200  loss:0.440077394247  mse:0.143687799573
  #     iteration:400  loss:0.440077394247  mse:0.143687844276
  #     iteration:600  loss:0.440077394247  mse:0.143687844276
  #     iteration:800  loss:0.440077424049  mse:0.143687844276
  #     iteration:999  loss:0.440077424049  mse:0.143687844276
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
      Default value: `tf.layers.dense`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "bernoulli").

  Returns:
    bernoulli: An instance of `tfd.Bernoulli`.
  """
  with tf.compat.v1.name_scope(name, 'bernoulli', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    logits = tf.squeeze(layer_fn(x, 1), axis=-1)
    return tfd.Bernoulli(logits=logits)

def bernoulli ( x , layer_fn = tf . compat . v1 . layers . dense , name = None ) : with tf . compat . v1 . name_scope ( name , 'bernoulli' , [ x ] ) : x = tf . convert_to_tensor ( value = x , name = 'x' ) logits = tf . squeeze ( layer_fn ( x , 1 ) , axis = - 1 ) return tfd . Bernoulli ( logits = logits )

Constructs a trainable tfd . Bernoulli distribution .


def normal(x,
           layer_fn=tf.compat.v1.layers.dense,
           loc_fn=lambda x: x,
           scale_fn=1.,
           name=None):
  """Constructs a trainable `tfd.Normal` distribution.

  This function creates a Normal distribution parameterized by loc and scale.
  Using default args, this function is mathematically equivalent to:

  ```none
  Y = Normal(loc=matmul(W, x) + b, scale=1)

  where,
    W in R^[d, n]
    b in R^d
  ```

  #### Examples

  This function can be used as a [linear regression](
  https://en.wikipedia.org/wiki/Linear_regression) loss.

  ```python
  # This example fits a linear regression loss.
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size).astype(dtype)
    b = np.random.randn(1).astype(dtype)
    true_mean = np.tensordot(x, w, axes=[[-1], [-1]]) + b
    noise = np.random.randn(n).astype(dtype)
    y = true_mean + noise
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting Normal maximum likelihood estimator.
  normal = tfp.trainable_distributions.normal(x)
  loss = -tf.reduce_mean(normal.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, normal.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:6.34114170074  mse:10.8444051743
  #     iteration:200  loss:1.40146839619  mse:0.965059816837
  #     iteration:400  loss:1.40052902699  mse:0.963181257248
  #     iteration:600  loss:1.40052902699  mse:0.963181257248
  #     iteration:800  loss:1.40052902699  mse:0.963181257248
  #     iteration:999  loss:1.40052902699  mse:0.963181257248
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
      Default value: `tf.layers.dense`.
    loc_fn: Python `callable` which transforms the `loc` parameter. Takes a
      (batch of) length-`dims` vectors and returns a `Tensor` of same shape and
      `dtype`.
      Default value: `lambda x: x`.
    scale_fn: Python `callable` or `Tensor`. If a `callable` transforms the
      `scale` parameters; if `Tensor` is the `tfd.Normal` `scale` argument.
      Takes a (batch of) length-`dims` vectors and returns a `Tensor` of same
      size. (Taking a `callable` or `Tensor` is how `tf.Variable` intializers
      behave.)
      Default value: `1`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "normal").

  Returns:
    normal: An instance of `tfd.Normal`.
  """
  with tf.compat.v1.name_scope(name, 'normal', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    if callable(scale_fn):
      y = layer_fn(x, 2)
      loc = loc_fn(y[..., 0])
      scale = scale_fn(y[..., 1])
    else:
      y = tf.squeeze(layer_fn(x, 1), axis=-1)
      loc = loc_fn(y)
      scale = tf.cast(scale_fn, loc.dtype.base_dtype)
    return tfd.Normal(loc=loc, scale=scale)

def normal ( x , layer_fn = tf . compat . v1 . layers . dense , loc_fn = lambda x : x , scale_fn = 1. , name = None ) : with tf . compat . v1 . name_scope ( name , 'normal' , [ x ] ) : x = tf . convert_to_tensor ( value = x , name = 'x' ) if callable ( scale_fn ) : y = layer_fn ( x , 2 ) loc = loc_fn ( y [ ... , 0 ] ) scale = scale_fn ( y [ ... , 1 ] ) else : y = tf . squeeze ( layer_fn ( x , 1 ) , axis = - 1 ) loc = loc_fn ( y ) scale = tf . cast ( scale_fn , loc . dtype . base_dtype ) return tfd . Normal ( loc = loc , scale = scale )

Constructs a trainable tfd . Normal distribution .


def poisson(x,
            layer_fn=tf.compat.v1.layers.dense,
            log_rate_fn=lambda x: x,
            name=None):
  """Constructs a trainable `tfd.Poisson` distribution.

  This function creates a Poisson distribution parameterized by log rate.
  Using default args, this function is mathematically equivalent to:

  ```none
  Y = Poisson(log_rate=matmul(W, x) + b)

  where,
    W in R^[d, n]
    b in R^d
  ```

  #### Examples

  This can be used as a [Poisson regression](
  https://en.wikipedia.org/wiki/Poisson_regression) loss.

  ```python
  # This example fits a poisson regression loss.
  import numpy as np
  import tensorflow as tf
  import tensorflow_probability as tfp

  # Create fictitious training data.
  dtype = np.float32
  n = 3000    # number of samples
  x_size = 4  # size of single x
  def make_training_data():
    np.random.seed(142)
    x = np.random.randn(n, x_size).astype(dtype)
    w = np.random.randn(x_size).astype(dtype)
    b = np.random.randn(1).astype(dtype)
    true_log_rate = np.tensordot(x, w, axes=[[-1], [-1]]) + b
    y = np.random.poisson(lam=np.exp(true_log_rate)).astype(dtype)
    return y, x
  y, x = make_training_data()

  # Build TF graph for fitting Poisson maximum likelihood estimator.
  poisson = tfp.trainable_distributions.poisson(x)
  loss = -tf.reduce_mean(poisson.log_prob(y))
  train_op = tf.train.AdamOptimizer(learning_rate=2.**-5).minimize(loss)
  mse = tf.reduce_mean(tf.squared_difference(y, poisson.mean()))
  init_op = tf.global_variables_initializer()

  # Run graph 1000 times.
  num_steps = 1000
  loss_ = np.zeros(num_steps)   # Style: `_` to indicate sess.run result.
  mse_ = np.zeros(num_steps)
  with tf.Session() as sess:
    sess.run(init_op)
    for it in xrange(loss_.size):
      _, loss_[it], mse_[it] = sess.run([train_op, loss, mse])
      if it % 200 == 0 or it == loss_.size - 1:
        print("iteration:{}  loss:{}  mse:{}".format(it, loss_[it], mse_[it]))

  # ==> iteration:0    loss:37.0814208984  mse:6359.41259766
  #     iteration:200  loss:1.42010736465  mse:40.7654914856
  #     iteration:400  loss:1.39027583599  mse:8.77660560608
  #     iteration:600  loss:1.3902695179   mse:8.78443241119
  #     iteration:800  loss:1.39026939869  mse:8.78443622589
  #     iteration:999  loss:1.39026939869  mse:8.78444766998
  ```

  Args:
    x: `Tensor` with floating type. Must have statically defined rank and
      statically known right-most dimension.
    layer_fn: Python `callable` which takes input `x` and `int` scalar `d` and
      returns a transformation of `x` with shape
      `tf.concat([tf.shape(x)[:-1], [1]], axis=0)`.
      Default value: `tf.layers.dense`.
    log_rate_fn: Python `callable` which transforms the `log_rate` parameter.
      Takes a (batch of) length-`dims` vectors and returns a `Tensor` of same
      shape and `dtype`.
      Default value: `lambda x: x`.
    name: A `name_scope` name for operations created by this function.
      Default value: `None` (i.e., "poisson").

  Returns:
    poisson: An instance of `tfd.Poisson`.
  """
  with tf.compat.v1.name_scope(name, 'poisson', [x]):
    x = tf.convert_to_tensor(value=x, name='x')
    log_rate = log_rate_fn(tf.squeeze(layer_fn(x, 1), axis=-1))
    return tfd.Poisson(log_rate=log_rate)

def poisson ( x , layer_fn = tf . compat . v1 . layers . dense , log_rate_fn = lambda x : x , name = None ) : with tf . compat . v1 . name_scope ( name , 'poisson' , [ x ] ) : x = tf . convert_to_tensor ( value = x , name = 'x' ) log_rate = log_rate_fn ( tf . squeeze ( layer_fn ( x , 1 ) , axis = - 1 ) ) return tfd . Poisson ( log_rate = log_rate )

Constructs a trainable tfd . Poisson distribution .


def _euler_method(random_draw_parts,
                  state_parts,
                  drift_parts,
                  step_size_parts,
                  volatility_parts,
                  name=None):
  """Applies one step of Euler-Maruyama method.

  Generates proposal of the form:

  ```python
  tfd.Normal(loc=state_parts + _get_drift(state_parts, ...),
             scale=tf.sqrt(step_size * volatility_fn(current_state)))
  ```

  `_get_drift(state_parts, ..)` is a diffusion drift value at `state_parts`.

  Args:
    random_draw_parts: Python `list` of `Tensor`s containing the value(s) of the
      random perturbation variable(s). Must broadcast with the shape of
      `state_parts`.
    state_parts: Python `list` of `Tensor`s representing the current
      state(s) of the Markov chain(s).
    drift_parts: Python `list` of `Tensor`s representing value of the drift
      `_get_drift(*state_parts, ..)`. Must broadcast with the shape of
      `state_parts`.
    step_size_parts: Python `list` of `Tensor`s representing the step size for
      the Euler-Maruyama method. Must broadcast with the shape of
      `state_parts`.  Larger step sizes lead to faster progress, but
      too-large step sizes make rejection exponentially more likely. When
      possible, it's often helpful to match per-variable step sizes to the
      standard deviations of the target distribution in each variable.
    volatility_parts: Python `list` of `Tensor`s representing the value of
      `volatility_fn(*state_parts)`. Must broadcast with the shape of
      `state_parts`.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'mala_euler_method').

  Returns:
    proposed_state_parts: Tensor or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at each result step. Has same shape as
      input `current_state_parts`.
  """
  with tf.compat.v1.name_scope(name, 'mala_euler_method', [
      random_draw_parts, state_parts, drift_parts, step_size_parts,
      volatility_parts
  ]):
    proposed_state_parts = []
    for random_draw, state, drift, step_size, volatility in zip(
        random_draw_parts,
        state_parts,
        drift_parts,
        step_size_parts,
        volatility_parts):
      proposal = state + drift + volatility * tf.sqrt(step_size) * random_draw
      proposed_state_parts.append(proposal)

    return proposed_state_parts

def _euler_method ( random_draw_parts , state_parts , drift_parts , step_size_parts , volatility_parts , name = None ) : with tf . compat . v1 . name_scope ( name , 'mala_euler_method' , [ random_draw_parts , state_parts , drift_parts , step_size_parts , volatility_parts ] ) : proposed_state_parts = [ ] for random_draw , state , drift , step_size , volatility in zip ( random_draw_parts , state_parts , drift_parts , step_size_parts , volatility_parts ) : proposal = state + drift + volatility * tf . sqrt ( step_size ) * random_draw proposed_state_parts . append ( proposal ) return proposed_state_parts

Applies one step of Euler - Maruyama method .


def _get_drift(step_size_parts, volatility_parts, grads_volatility,
               grads_target_log_prob,
               name=None):
  """Compute diffusion drift at the current location `current_state`.

  The drift of the diffusion at is computed as

  ```none
  0.5 * `step_size` * volatility_parts * `target_log_prob_fn(current_state)`
  + `step_size` * `grads_volatility`
  ```

  where `volatility_parts` = `volatility_fn(current_state)**2` and
  `grads_volatility` is a gradient of `volatility_parts` at the `current_state`.

  Args:
    step_size_parts: Python `list` of `Tensor`s representing the step size for
      Euler-Maruyama method. Must broadcast with the shape of
      `volatility_parts`.  Larger step sizes lead to faster progress, but
      too-large step sizes make rejection exponentially more likely. When
      possible, it's often helpful to match per-variable step sizes to the
      standard deviations of the target distribution in each variable.
    volatility_parts: Python `list` of `Tensor`s representing the value of
      `volatility_fn(*state_parts)`.
    grads_volatility: Python list of `Tensor`s representing the value of the
      gradient of `volatility_parts**2` wrt the state of the chain.
    grads_target_log_prob: Python list of `Tensor`s representing
      gradient of `target_log_prob_fn(*state_parts`) wrt `state_parts`. Must
      have same shape as `volatility_parts`.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'mala_get_drift').

  Returns:
    drift_parts: Tensor or Python list of `Tensor`s representing the
      state(s) of the Markov chain(s) at each result step. Has same shape as
      input `current_state_parts`.
  """

  with tf.compat.v1.name_scope(name, 'mala_get_drift', [
      step_size_parts, volatility_parts, grads_volatility, grads_target_log_prob
  ]):

    drift_parts = []

    for step_size, volatility, grad_volatility, grad_target_log_prob in (
        zip(step_size_parts,
            volatility_parts,
            grads_volatility,
            grads_target_log_prob)):
      volatility_squared = tf.square(volatility)
      drift = 0.5 * step_size * (volatility_squared * grad_target_log_prob
                                 + grad_volatility)
      drift_parts.append(drift)

    return drift_parts

def _get_drift ( step_size_parts , volatility_parts , grads_volatility , grads_target_log_prob , name = None ) : with tf . compat . v1 . name_scope ( name , 'mala_get_drift' , [ step_size_parts , volatility_parts , grads_volatility , grads_target_log_prob ] ) : drift_parts = [ ] for step_size , volatility , grad_volatility , grad_target_log_prob in ( zip ( step_size_parts , volatility_parts , grads_volatility , grads_target_log_prob ) ) : volatility_squared = tf . square ( volatility ) drift = 0.5 * step_size * ( volatility_squared * grad_target_log_prob + grad_volatility ) drift_parts . append ( drift ) return drift_parts

Compute diffusion drift at the current location current_state .


def _compute_log_acceptance_correction(current_state_parts,
                                       proposed_state_parts,
                                       current_volatility_parts,
                                       proposed_volatility_parts,
                                       current_drift_parts,
                                       proposed_drift_parts,
                                       step_size_parts,
                                       independent_chain_ndims,
                                       name=None):
  r"""Helper to `kernel` which computes the log acceptance-correction.

  Computes `log_acceptance_correction` as described in `MetropolisHastings`
  class. The proposal density is normal. More specifically,

   ```none
  q(proposed_state | current_state) \sim N(current_state + current_drift,
  step_size * current_volatility**2)

  q(current_state | proposed_state) \sim N(proposed_state + proposed_drift,
  step_size * proposed_volatility**2)
  ```

  The `log_acceptance_correction` is then

  ```none
  log_acceptance_correctio = q(current_state | proposed_state)
  - q(proposed_state | current_state)
  ```

  Args:
    current_state_parts: Python `list` of `Tensor`s representing the value(s) of
      the current state of the chain.
    proposed_state_parts:  Python `list` of `Tensor`s representing the value(s)
      of the proposed state of the chain. Must broadcast with the shape of
      `current_state_parts`.
    current_volatility_parts: Python `list` of `Tensor`s representing the value
      of `volatility_fn(*current_volatility_parts)`. Must broadcast with the
      shape of `current_state_parts`.
    proposed_volatility_parts: Python `list` of `Tensor`s representing the value
      of `volatility_fn(*proposed_volatility_parts)`. Must broadcast with the
      shape of `current_state_parts`
    current_drift_parts: Python `list` of `Tensor`s representing value of the
      drift `_get_drift(*current_state_parts, ..)`. Must broadcast with the
      shape of `current_state_parts`.
    proposed_drift_parts: Python `list` of `Tensor`s representing value of the
      drift `_get_drift(*proposed_drift_parts, ..)`. Must broadcast with the
      shape of `current_state_parts`.
    step_size_parts: Python `list` of `Tensor`s representing the step size for
      Euler-Maruyama method. Must broadcast with the shape of
      `current_state_parts`.
    independent_chain_ndims: Scalar `int` `Tensor` representing the number of
      leftmost `Tensor` dimensions which index independent chains.
    name: Python `str` name prefixed to Ops created by this function.
      Default value: `None` (i.e., 'compute_log_acceptance_correction').

  Returns:
    log_acceptance_correction: `Tensor` representing the `log`
      acceptance-correction.  (See docstring for mathematical definition.)
  """

  with tf.compat.v1.name_scope(name, 'compute_log_acceptance_correction', [
      current_state_parts, proposed_state_parts, current_volatility_parts,
      proposed_volatility_parts, current_drift_parts, proposed_drift_parts,
      step_size_parts, independent_chain_ndims
  ]):

    proposed_log_density_parts = []
    dual_log_density_parts = []

    for [
        current_state,
        proposed_state,
        current_volatility,
        proposed_volatility,
        current_drift,
        proposed_drift,
        step_size,
    ] in zip(
        current_state_parts,
        proposed_state_parts,
        current_volatility_parts,
        proposed_volatility_parts,
        current_drift_parts,
        proposed_drift_parts,
        step_size_parts,
    ):
      axis = tf.range(independent_chain_ndims, tf.rank(current_state))

      state_diff = proposed_state - current_state

      current_volatility *= tf.sqrt(step_size)

      proposed_energy = (state_diff - current_drift) / current_volatility

      proposed_volatility *= tf.sqrt(step_size)
      # Compute part of `q(proposed_state | current_state)`
      proposed_energy = (
          tf.reduce_sum(
              input_tensor=mcmc_util.safe_sum(
                  [tf.math.log(current_volatility),
                   0.5 * (proposed_energy**2)]),
              axis=axis))
      proposed_log_density_parts.append(-proposed_energy)

      # Compute part of `q(current_state | proposed_state)`
      dual_energy = (state_diff + proposed_drift) / proposed_volatility
      dual_energy = (
          tf.reduce_sum(
              input_tensor=mcmc_util.safe_sum(
                  [tf.math.log(proposed_volatility), 0.5 * (dual_energy**2)]),
              axis=axis))
      dual_log_density_parts.append(-dual_energy)

    # Compute `q(proposed_state | current_state)`
    proposed_log_density_reduce = tf.reduce_sum(
        input_tensor=tf.stack(proposed_log_density_parts, axis=-1), axis=-1)
    # Compute `q(current_state | proposed_state)`
    dual_log_density_reduce = tf.reduce_sum(
        input_tensor=tf.stack(dual_log_density_parts, axis=-1), axis=-1)

    return mcmc_util.safe_sum([dual_log_density_reduce,
                               -proposed_log_density_reduce])

def _compute_log_acceptance_correction ( current_state_parts , proposed_state_parts , current_volatility_parts , proposed_volatility_parts , current_drift_parts , proposed_drift_parts , step_size_parts , independent_chain_ndims , name = None ) : r with tf . compat . v1 . name_scope ( name , 'compute_log_acceptance_correction' , [ current_state_parts , proposed_state_parts , current_volatility_parts , proposed_volatility_parts , current_drift_parts , proposed_drift_parts , step_size_parts , independent_chain_ndims ] ) : proposed_log_density_parts = [ ] dual_log_density_parts = [ ] for [ current_state , proposed_state , current_volatility , proposed_volatility , current_drift , proposed_drift , step_size , ] in zip ( current_state_parts , proposed_state_parts , current_volatility_parts , proposed_volatility_parts , current_drift_parts , proposed_drift_parts , step_size_parts , ) : axis = tf . range ( independent_chain_ndims , tf . rank ( current_state ) ) state_diff = proposed_state - current_state current_volatility *= tf . sqrt ( step_size ) proposed_energy = ( state_diff - current_drift ) / current_volatility proposed_volatility *= tf . sqrt ( step_size ) proposed_energy = ( tf . reduce_sum ( input_tensor = mcmc_util . safe_sum ( [ tf . math . log ( current_volatility ) , 0.5 * ( proposed_energy ** 2 ) ] ) , axis = axis ) ) proposed_log_density_parts . append ( - proposed_energy ) dual_energy = ( state_diff + proposed_drift ) / proposed_volatility dual_energy = ( tf . reduce_sum ( input_tensor = mcmc_util . safe_sum ( [ tf . math . log ( proposed_volatility ) , 0.5 * ( dual_energy ** 2 ) ] ) , axis = axis ) ) dual_log_density_parts . append ( - dual_energy ) proposed_log_density_reduce = tf . reduce_sum ( input_tensor = tf . stack ( proposed_log_density_parts , axis = - 1 ) , axis = - 1 ) dual_log_density_reduce = tf . reduce_sum ( input_tensor = tf . stack ( dual_log_density_parts , axis = - 1 ) , axis = - 1 ) return mcmc_util . safe_sum ( [ dual_log_density_reduce , - proposed_log_density_reduce ] )

r Helper to kernel which computes the log acceptance - correction .


def _maybe_call_volatility_fn_and_grads(volatility_fn,
                                        state,
                                        volatility_fn_results=None,
                                        grads_volatility_fn=None,
                                        sample_shape=None,
                                        parallel_iterations=10):
  """Helper which computes `volatility_fn` results and grads, if needed."""
  state_parts = list(state) if mcmc_util.is_list_like(state) else [state]
  needs_volatility_fn_gradients = grads_volatility_fn is None

  # Convert `volatility_fn_results` to a list
  if volatility_fn_results is None:
    volatility_fn_results = volatility_fn(*state_parts)

  volatility_fn_results = (list(volatility_fn_results)
                           if mcmc_util.is_list_like(volatility_fn_results)
                           else [volatility_fn_results])
  if len(volatility_fn_results) == 1:
    volatility_fn_results *= len(state_parts)
  if len(state_parts) != len(volatility_fn_results):
    raise ValueError('`volatility_fn` should return a tensor or a list '
                     'of the same length as `current_state`.')

  # The shape of 'volatility_parts' needs to have the number of chains as a
  # leading dimension. For determinism we broadcast 'volatility_parts' to the
  # shape of `state_parts` since each dimension of `state_parts` could have a
  # different volatility value.

  volatility_fn_results = _maybe_broadcast_volatility(volatility_fn_results,
                                                      state_parts)
  if grads_volatility_fn is None:
    [
        _,
        grads_volatility_fn,
    ] = diag_jacobian(
        xs=state_parts,
        ys=volatility_fn_results,
        sample_shape=sample_shape,
        parallel_iterations=parallel_iterations,
        fn=volatility_fn)

  # Compute gradient of `volatility_parts**2`
  if needs_volatility_fn_gradients:
    grads_volatility_fn = [
        2. * g * volatility if g is not None else tf.zeros_like(
            fn_arg, dtype=fn_arg.dtype.base_dtype)
        for g, volatility, fn_arg in zip(
            grads_volatility_fn, volatility_fn_results, state_parts)
    ]

  return volatility_fn_results, grads_volatility_fn

Helper which computes volatility_fn results and grads if needed .


def _maybe_broadcast_volatility(volatility_parts,
                                state_parts):
  """Helper to broadcast `volatility_parts` to the shape of `state_parts`."""
  return [v + tf.zeros_like(sp, dtype=sp.dtype.base_dtype)
          for v, sp in zip(volatility_parts, state_parts)]

def _maybe_broadcast_volatility ( volatility_parts , state_parts ) : return [ v + tf . zeros_like ( sp , dtype = sp . dtype . base_dtype ) for v , sp in zip ( volatility_parts , state_parts ) ]

Helper to broadcast volatility_parts to the shape of state_parts .


def make_ar_transition_matrix(coefficients):
  """Build transition matrix for an autoregressive StateSpaceModel.

  When applied to a vector of previous values, this matrix computes
  the expected new value (summing the previous states according to the
  autoregressive coefficients) in the top dimension of the state space,
  and moves all previous values down by one dimension, 'forgetting' the
  final (least recent) value. That is, it looks like this:

  ```
  ar_matrix = [ coefs[0], coefs[1], ..., coefs[order]
                1.,       0 ,       ..., 0.
                0.,       1.,       ..., 0.
                ...
                0.,       0.,  ..., 1.,  0.            ]
  ```

  Args:
    coefficients: float `Tensor` of shape `concat([batch_shape, [order]])`.

  Returns:
    ar_matrix: float `Tensor` with shape `concat([batch_shape,
    [order, order]])`.
  """

  top_row = tf.expand_dims(coefficients, -2)
  coef_shape = dist_util.prefer_static_shape(coefficients)
  batch_shape, order = coef_shape[:-1], coef_shape[-1]
  remaining_rows = tf.concat([
      tf.eye(order - 1, dtype=coefficients.dtype, batch_shape=batch_shape),
      tf.zeros(tf.concat([batch_shape, (order - 1, 1)], axis=0),
               dtype=coefficients.dtype)
  ], axis=-1)
  ar_matrix = tf.concat([top_row, remaining_rows], axis=-2)
  return ar_matrix

def make_ar_transition_matrix ( coefficients ) : top_row = tf . expand_dims ( coefficients , - 2 ) coef_shape = dist_util . prefer_static_shape ( coefficients ) batch_shape , order = coef_shape [ : - 1 ] , coef_shape [ - 1 ] remaining_rows = tf . concat ( [ tf . eye ( order - 1 , dtype = coefficients . dtype , batch_shape = batch_shape ) , tf . zeros ( tf . concat ( [ batch_shape , ( order - 1 , 1 ) ] , axis = 0 ) , dtype = coefficients . dtype ) ] , axis = - 1 ) ar_matrix = tf . concat ( [ top_row , remaining_rows ] , axis = - 2 ) return ar_matrix

Build transition matrix for an autoregressive StateSpaceModel .


def _sample_shape(self, x):
    """Computes graph and static `sample_shape`."""
    x_ndims = (
        tf.rank(x) if tensorshape_util.rank(x.shape) is None else
        tensorshape_util.rank(x.shape))
    event_ndims = (
        tf.size(input=self.event_shape_tensor())
        if tensorshape_util.rank(self.event_shape) is None else
        tensorshape_util.rank(self.event_shape))
    batch_ndims = (
        tf.size(input=self._batch_shape_unexpanded)
        if tensorshape_util.rank(self.batch_shape) is None else
        tensorshape_util.rank(self.batch_shape))
    sample_ndims = x_ndims - batch_ndims - event_ndims
    if isinstance(sample_ndims, int):
      static_sample_shape = x.shape[:sample_ndims]
    else:
      static_sample_shape = tf.TensorShape(None)
    if tensorshape_util.is_fully_defined(static_sample_shape):
      sample_shape = np.int32(static_sample_shape)
    else:
      sample_shape = tf.shape(input=x)[:sample_ndims]
    return sample_shape, static_sample_shape

def _sample_shape ( self , x ) : x_ndims = ( tf . rank ( x ) if tensorshape_util . rank ( x . shape ) is None else tensorshape_util . rank ( x . shape ) ) event_ndims = ( tf . size ( input = self . event_shape_tensor ( ) ) if tensorshape_util . rank ( self . event_shape ) is None else tensorshape_util . rank ( self . event_shape ) ) batch_ndims = ( tf . size ( input = self . _batch_shape_unexpanded ) if tensorshape_util . rank ( self . batch_shape ) is None else tensorshape_util . rank ( self . batch_shape ) ) sample_ndims = x_ndims - batch_ndims - event_ndims if isinstance ( sample_ndims , int ) : static_sample_shape = x . shape [ : sample_ndims ] else : static_sample_shape = tf . TensorShape ( None ) if tensorshape_util . is_fully_defined ( static_sample_shape ) : sample_shape = np . int32 ( static_sample_shape ) else : sample_shape = tf . shape ( input = x ) [ : sample_ndims ] return sample_shape , static_sample_shape

Computes graph and static sample_shape .


def _call_reshape_input_output(self, fn, x, extra_kwargs=None):
    """Calls `fn`, appropriately reshaping its input `x` and output."""
    # Note: we take `extra_kwargs` as a dict rather than `**extra_kwargs`
    # because it is possible the user provided extra kwargs would itself
    # have `fn` and/or `x` as a key.
    with tf.control_dependencies(self._runtime_assertions +
                                 self._validate_sample_arg(x)):
      sample_shape, static_sample_shape = self._sample_shape(x)
      old_shape = tf.concat(
          [
              sample_shape,
              self.distribution.batch_shape_tensor(),
              self.event_shape_tensor(),
          ],
          axis=0)
      x_reshape = tf.reshape(x, old_shape)
      result = fn(x_reshape, **extra_kwargs) if extra_kwargs else fn(x_reshape)
      new_shape = tf.concat(
          [
              sample_shape,
              self._batch_shape_unexpanded,
          ], axis=0)
      result = tf.reshape(result, new_shape)
      if (tensorshape_util.rank(static_sample_shape) is not None and
          tensorshape_util.rank(self.batch_shape) is not None):
        new_shape = tensorshape_util.concatenate(static_sample_shape,
                                                 self.batch_shape)
        tensorshape_util.set_shape(result, new_shape)
      return result

def _call_reshape_input_output ( self , fn , x , extra_kwargs = None ) : with tf . control_dependencies ( self . _runtime_assertions + self . _validate_sample_arg ( x ) ) : sample_shape , static_sample_shape = self . _sample_shape ( x ) old_shape = tf . concat ( [ sample_shape , self . distribution . batch_shape_tensor ( ) , self . event_shape_tensor ( ) , ] , axis = 0 ) x_reshape = tf . reshape ( x , old_shape ) result = fn ( x_reshape , ** extra_kwargs ) if extra_kwargs else fn ( x_reshape ) new_shape = tf . concat ( [ sample_shape , self . _batch_shape_unexpanded , ] , axis = 0 ) result = tf . reshape ( result , new_shape ) if ( tensorshape_util . rank ( static_sample_shape ) is not None and tensorshape_util . rank ( self . batch_shape ) is not None ) : new_shape = tensorshape_util . concatenate ( static_sample_shape , self . batch_shape ) tensorshape_util . set_shape ( result , new_shape ) return result

Calls fn appropriately reshaping its input x and output .


def _call_and_reshape_output(
      self,
      fn,
      event_shape_list=None,
      static_event_shape_list=None,
      extra_kwargs=None):
    """Calls `fn` and appropriately reshapes its output."""
    # Note: we take `extra_kwargs` as a dict rather than `**extra_kwargs`
    # because it is possible the user provided extra kwargs would itself
    # have `fn`, `event_shape_list`, `static_event_shape_list` and/or
    # `extra_kwargs` as keys.
    with tf.control_dependencies(self._runtime_assertions):
      if event_shape_list is None:
        event_shape_list = [self._event_shape_tensor()]
      if static_event_shape_list is None:
        static_event_shape_list = [self.event_shape]
      new_shape = tf.concat(
          [self._batch_shape_unexpanded] + event_shape_list, axis=0)
      result = tf.reshape(fn(**extra_kwargs) if extra_kwargs else fn(),
                          new_shape)
      if (tensorshape_util.rank(self.batch_shape) is not None and
          tensorshape_util.rank(self.event_shape) is not None):
        event_shape = tf.TensorShape([])
        for rss in static_event_shape_list:
          event_shape = tensorshape_util.concatenate(event_shape, rss)
        static_shape = tensorshape_util.concatenate(
            self.batch_shape, event_shape)
        tensorshape_util.set_shape(result, static_shape)
      return result

def _call_and_reshape_output ( self , fn , event_shape_list = None , static_event_shape_list = None , extra_kwargs = None ) : with tf . control_dependencies ( self . _runtime_assertions ) : if event_shape_list is None : event_shape_list = [ self . _event_shape_tensor ( ) ] if static_event_shape_list is None : static_event_shape_list = [ self . event_shape ] new_shape = tf . concat ( [ self . _batch_shape_unexpanded ] + event_shape_list , axis = 0 ) result = tf . reshape ( fn ( ** extra_kwargs ) if extra_kwargs else fn ( ) , new_shape ) if ( tensorshape_util . rank ( self . batch_shape ) is not None and tensorshape_util . rank ( self . event_shape ) is not None ) : event_shape = tf . TensorShape ( [ ] ) for rss in static_event_shape_list : event_shape = tensorshape_util . concatenate ( event_shape , rss ) static_shape = tensorshape_util . concatenate ( self . batch_shape , event_shape ) tensorshape_util . set_shape ( result , static_shape ) return result

Calls fn and appropriately reshapes its output .


def _bdtr(k, n, p):
  """The binomial cumulative distribution function.

  Args:
    k: floating point `Tensor`.
    n: floating point `Tensor`.
    p: floating point `Tensor`.

  Returns:
    `sum_{j=0}^k p^j (1 - p)^(n - j)`.
  """
  # Trick for getting safe backprop/gradients into n, k when
  #   betainc(a = 0, ..) = nan
  # Write:
  #   where(unsafe, safe_output, betainc(where(unsafe, safe_input, input)))
  ones = tf.ones_like(n - k)
  k_eq_n = tf.equal(k, n)
  safe_dn = tf.where(k_eq_n, ones, n - k)
  dk = tf.math.betainc(a=safe_dn, b=k + 1, x=1 - p)
  return tf.where(k_eq_n, ones, dk)

def _bdtr ( k , n , p ) : ones = tf . ones_like ( n - k ) k_eq_n = tf . equal ( k , n ) safe_dn = tf . where ( k_eq_n , ones , n - k ) dk = tf . math . betainc ( a = safe_dn , b = k + 1 , x = 1 - p ) return tf . where ( k_eq_n , ones , dk )

The binomial cumulative distribution function .


def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):
    """Executes `model`, creating both samples and distributions."""
    ds = []
    values_out = []
    seed = seed_stream.SeedStream('JointDistributionCoroutine', seed)
    gen = self._model()
    index = 0
    d = next(gen)
    try:
      while True:
        actual_distribution = d.distribution if isinstance(d, self.Root) else d
        ds.append(actual_distribution)
        if (value is not None and len(value) > index and
            value[index] is not None):
          seed()
          next_value = value[index]
        else:
          next_value = actual_distribution.sample(
              sample_shape=sample_shape if isinstance(d, self.Root) else (),
              seed=seed())
        values_out.append(next_value)
        index += 1
        d = gen.send(next_value)
    except StopIteration:
      pass
    return ds, values_out

def _flat_sample_distributions ( self , sample_shape = ( ) , seed = None , value = None ) : ds = [ ] values_out = [ ] seed = seed_stream . SeedStream ( 'JointDistributionCoroutine' , seed ) gen = self . _model ( ) index = 0 d = next ( gen ) try : while True : actual_distribution = d . distribution if isinstance ( d , self . Root ) else d ds . append ( actual_distribution ) if ( value is not None and len ( value ) > index and value [ index ] is not None ) : seed ( ) next_value = value [ index ] else : next_value = actual_distribution . sample ( sample_shape = sample_shape if isinstance ( d , self . Root ) else ( ) , seed = seed ( ) ) values_out . append ( next_value ) index += 1 d = gen . send ( next_value ) except StopIteration : pass return ds , values_out

Executes model creating both samples and distributions .


def latent_dirichlet_allocation(concentration, topics_words):
  """Latent Dirichlet Allocation in terms of its generative process.

  The model posits a distribution over bags of words and is parameterized by
  a concentration and the topic-word probabilities. It collapses per-word
  topic assignments.

  Args:
    concentration: A Tensor of shape [1, num_topics], which parameterizes the
      Dirichlet prior over topics.
    topics_words: A Tensor of shape [num_topics, num_words], where each row
      (topic) denotes the probability of each word being in that topic.

  Returns:
    bag_of_words: A random variable capturing a sample from the model, of shape
      [1, num_words]. It represents one generated document as a bag of words.
  """
  topics = ed.Dirichlet(concentration=concentration, name="topics")
  word_probs = tf.matmul(topics, topics_words)
  # The observations are bags of words and therefore not one-hot. However,
  # log_prob of OneHotCategorical computes the probability correctly in
  # this case.
  bag_of_words = ed.OneHotCategorical(probs=word_probs, name="bag_of_words")
  return bag_of_words

def latent_dirichlet_allocation ( concentration , topics_words ) : topics = ed . Dirichlet ( concentration = concentration , name = "topics" ) word_probs = tf . matmul ( topics , topics_words ) bag_of_words = ed . OneHotCategorical ( probs = word_probs , name = "bag_of_words" ) return bag_of_words

Latent Dirichlet Allocation in terms of its generative process .


def make_lda_variational(activation, num_topics, layer_sizes):
  """Creates the variational distribution for LDA.

  Args:
    activation: Activation function to use.
    num_topics: The number of topics.
    layer_sizes: The number of hidden units per layer in the encoder.

  Returns:
    lda_variational: A function that takes a bag-of-words Tensor as
      input and returns a distribution over topics.
  """
  encoder_net = tf.keras.Sequential()
  for num_hidden_units in layer_sizes:
    encoder_net.add(
        tf.keras.layers.Dense(
            num_hidden_units,
            activation=activation,
            kernel_initializer=tf.compat.v1.glorot_normal_initializer()))
  encoder_net.add(
      tf.keras.layers.Dense(
          num_topics,
          activation=tf.nn.softplus,
          kernel_initializer=tf.compat.v1.glorot_normal_initializer()))

  def lda_variational(bag_of_words):
    concentration = _clip_dirichlet_parameters(encoder_net(bag_of_words))
    return ed.Dirichlet(concentration=concentration, name="topics_posterior")

  return lda_variational

def make_lda_variational ( activation , num_topics , layer_sizes ) : encoder_net = tf . keras . Sequential ( ) for num_hidden_units in layer_sizes : encoder_net . add ( tf . keras . layers . Dense ( num_hidden_units , activation = activation , kernel_initializer = tf . compat . v1 . glorot_normal_initializer ( ) ) ) encoder_net . add ( tf . keras . layers . Dense ( num_topics , activation = tf . nn . softplus , kernel_initializer = tf . compat . v1 . glorot_normal_initializer ( ) ) ) def lda_variational ( bag_of_words ) : concentration = _clip_dirichlet_parameters ( encoder_net ( bag_of_words ) ) return ed . Dirichlet ( concentration = concentration , name = "topics_posterior" ) return lda_variational

Creates the variational distribution for LDA .


def get_topics_strings(topics_words, alpha, vocabulary,
                       topics_to_print=10, words_per_topic=10):
  """Returns the summary of the learned topics.

  Arguments:
    topics_words: KxV tensor with topics as rows and words as columns.
    alpha: 1xK tensor of prior Dirichlet concentrations for the
        topics.
    vocabulary: A mapping of word's integer index to the corresponding string.
    topics_to_print: The number of topics with highest prior weight to
        summarize.
    words_per_topic: Number of wodrs per topic to return.

  Returns:
    summary: A np.array with strings.
  """
  alpha = np.squeeze(alpha, axis=0)
  # Use a stable sorting algorithm so that when alpha is fixed
  # we always get the same topics.
  highest_weight_topics = np.argsort(-alpha, kind="mergesort")
  top_words = np.argsort(-topics_words, axis=1)

  res = []
  for topic_idx in highest_weight_topics[:topics_to_print]:
    l = ["index={} alpha={:.2f}".format(topic_idx, alpha[topic_idx])]
    l += [vocabulary[word] for word in top_words[topic_idx, :words_per_topic]]
    res.append(" ".join(l))

  return np.array(res)

def
get_topics_strings
(
topics_words
,
alpha
,
vocabulary
,
topics_to_print
=
10
,
words_per_topic
=
10
)
:
alpha
=
np
.
squeeze
(
alpha
,
axis
=
0
)
highest_weight_topics
=
np
.
argsort
(
alpha
,
kind
=
"mergesort"
)
top_words
=
np
.
argsort
(
topics_words
,
axis
=
1
)
res
=
[
]
for
topic_idx
in
highest_weight_topics
[
:
topics_to_print
]
:
l
=
[
"index={} alpha={:.2f}"
.
format
(
topic_idx
,
alpha
[
topic_idx
]
)
]
l
+=
[
vocabulary
[
word
]
for
word
in
top_words
[
topic_idx
,
:
words_per_topic
]
]
res
.
append
(
" "
.
join
(
l
)
)
return
np
.
array
(
res
)

Returns the summary of the learned topics .


def newsgroups_dataset(directory, split_name, num_words, shuffle_and_repeat):
  """20 newsgroups as a tf.data.Dataset."""
  data = np.load(download(directory, FILE_TEMPLATE.format(split=split_name)))
  # The last row is empty in both train and test.
  data = data[:-1]

  # Each row is a list of word ids in the document. We first convert this to
  # sparse COO matrix (which automatically sums the repeating words). Then,
  # we convert this COO matrix to CSR format which allows for fast querying of
  # documents.
  num_documents = data.shape[0]
  indices = np.array([(row_idx, column_idx)
                      for row_idx, row in enumerate(data)
                      for column_idx in row])
  sparse_matrix = scipy.sparse.coo_matrix(
      (np.ones(indices.shape[0]), (indices[:, 0], indices[:, 1])),
      shape=(num_documents, num_words),
      dtype=np.float32)
  sparse_matrix = sparse_matrix.tocsr()

  dataset = tf.data.Dataset.range(num_documents)

  # For training, we shuffle each epoch and repeat the epochs.
  if shuffle_and_repeat:
    dataset = dataset.shuffle(num_documents).repeat()

  # Returns a single document as a dense TensorFlow tensor. The dataset is
  # stored as a sparse matrix outside of the graph.
  def get_row_py_func(idx):
    def get_row_python(idx_py):
      return np.squeeze(np.array(sparse_matrix[idx_py].todense()), axis=0)

    py_func = tf.compat.v1.py_func(
        get_row_python, [idx], tf.float32, stateful=False)
    py_func.set_shape((num_words,))
    return py_func

  dataset = dataset.map(get_row_py_func)
  return dataset

def newsgroups_dataset ( directory , split_name , num_words , shuffle_and_repeat ) : data = np . load ( download ( directory , FILE_TEMPLATE . format ( split = split_name ) ) ) data = data [ : - 1 ] num_documents = data . shape [ 0 ] indices = np . array ( [ ( row_idx , column_idx ) for row_idx , row in enumerate ( data ) for column_idx in row ] ) sparse_matrix = scipy . sparse . coo_matrix ( ( np . ones ( indices . shape [ 0 ] ) , ( indices [ : , 0 ] , indices [ : , 1 ] ) ) , shape = ( num_documents , num_words ) , dtype = np . float32 ) sparse_matrix = sparse_matrix . tocsr ( ) dataset = tf . data . Dataset . range ( num_documents ) if shuffle_and_repeat : dataset = dataset . shuffle ( num_documents ) . repeat ( ) def get_row_py_func ( idx ) : def get_row_python ( idx_py ) : return np . squeeze ( np . array ( sparse_matrix [ idx_py ] . todense ( ) ) , axis = 0 ) py_func = tf . compat . v1 . py_func ( get_row_python , [ idx ] , tf . float32 , stateful = False ) py_func . set_shape ( ( num_words , ) ) return py_func dataset = dataset . map ( get_row_py_func ) return dataset

20 newsgroups as a tf . data . Dataset .


def build_fake_input_fns(batch_size):
  """Builds fake data for unit testing."""
  num_words = 1000
  vocabulary = [str(i) for i in range(num_words)]

  random_sample = np.random.randint(
      10, size=(batch_size, num_words)).astype(np.float32)

  def train_input_fn():
    dataset = tf.data.Dataset.from_tensor_slices(random_sample)
    dataset = dataset.batch(batch_size).repeat()
    return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

  def eval_input_fn():
    dataset = tf.data.Dataset.from_tensor_slices(random_sample)
    dataset = dataset.batch(batch_size)
    return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

  return train_input_fn, eval_input_fn, vocabulary

def build_fake_input_fns ( batch_size ) : num_words = 1000 vocabulary = [ str ( i ) for i in range ( num_words ) ] random_sample = np . random . randint ( 10 , size = ( batch_size , num_words ) ) . astype ( np . float32 ) def train_input_fn ( ) : dataset = tf . data . Dataset . from_tensor_slices ( random_sample ) dataset = dataset . batch ( batch_size ) . repeat ( ) return tf . compat . v1 . data . make_one_shot_iterator ( dataset ) . get_next ( ) def eval_input_fn ( ) : dataset = tf . data . Dataset . from_tensor_slices ( random_sample ) dataset = dataset . batch ( batch_size ) return tf . compat . v1 . data . make_one_shot_iterator ( dataset ) . get_next ( ) return train_input_fn , eval_input_fn , vocabulary


def build_input_fns(data_dir, batch_size):
  """Builds iterators for train and evaluation data.

  Each object is represented as a bag-of-words vector.

  Arguments:
    data_dir: Folder in which to store the data.
    batch_size: Batch size for both train and evaluation.

  Returns:
    train_input_fn: A function that returns an iterator over the training data.
    eval_input_fn: A function that returns an iterator over the evaluation data.
    vocabulary: A mapping of word's integer index to the corresponding string.
  """

  with open(download(data_dir, "vocab.pkl"), "r") as f:
    words_to_idx = pickle.load(f)
  num_words = len(words_to_idx)

  vocabulary = [None] * num_words
  for word, idx in words_to_idx.items():
    vocabulary[idx] = word

  # Build an iterator over training batches.
  def train_input_fn():
    dataset = newsgroups_dataset(
        data_dir, "train", num_words, shuffle_and_repeat=True)
    # Prefetching makes training about 1.5x faster.
    dataset = dataset.batch(batch_size).prefetch(32)
    return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

  # Build an iterator over the heldout set.
  def eval_input_fn():
    dataset = newsgroups_dataset(
        data_dir, "test", num_words, shuffle_and_repeat=False)
    dataset = dataset.batch(batch_size)
    return tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()

  return train_input_fn, eval_input_fn, vocabulary

def build_input_fns ( data_dir , batch_size ) : with open ( download ( data_dir , "vocab.pkl" ) , "r" ) as f : words_to_idx = pickle . load ( f ) num_words = len ( words_to_idx ) vocabulary = [ None ] * num_words for word , idx in words_to_idx . items ( ) : vocabulary [ idx ] = word def train_input_fn ( ) : dataset = newsgroups_dataset ( data_dir , "train" , num_words , shuffle_and_repeat = True ) dataset = dataset . batch ( batch_size ) . prefetch ( 32 ) return tf . compat . v1 . data . make_one_shot_iterator ( dataset ) . get_next ( ) def eval_input_fn ( ) : dataset = newsgroups_dataset ( data_dir , "test" , num_words , shuffle_and_repeat = False ) dataset = dataset . batch ( batch_size ) return tf . compat . v1 . data . make_one_shot_iterator ( dataset ) . get_next ( ) return train_input_fn , eval_input_fn , vocabulary

Builds iterators for train and evaluation data .