if not tf.executing_eagerly():<EOL><INDENT>raise NotImplementedError("<STR_LIT>")<EOL><DEDENT>with tf.compat.v1.name_scope(<EOL>name,<EOL>default_name="<STR_LIT>",<EOL>values=[<EOL>current_state, step_size, seed, current_target_log_prob,<EOL>current_grads_target_log_prob<EOL>]):<EOL><INDENT>with tf.compat.v1.name_scope("<STR_LIT>"):<EOL><INDENT>current_state = [tf.convert_to_tensor(value=s) for s in current_state]<EOL>step_size = [tf.convert_to_tensor(value=s) for s in step_size]<EOL>value_and_gradients_fn = lambda *args: tfp.math.value_and_gradient(  <EOL>target_log_prob_fn, args)<EOL>value_and_gradients_fn = _embed_no_none_gradient_check(<EOL>value_and_gradients_fn)<EOL>if (current_target_log_prob is None or<EOL>current_grads_target_log_prob is None):<EOL><INDENT>(current_target_log_prob,<EOL>current_grads_target_log_prob) = value_and_gradients_fn(*current_state)<EOL><DEDENT>seed_stream = tfd.SeedStream(seed, "<STR_LIT>")<EOL>current_momentum = []<EOL>for state_tensor in current_state:<EOL><INDENT>momentum_tensor = tf.random.normal(<EOL>shape=tf.shape(input=state_tensor),<EOL>dtype=state_tensor.dtype,<EOL>seed=seed_stream())<EOL>current_momentum.append(momentum_tensor)<EOL><DEDENT>log_slice_sample = tf.math.log(tf.random.uniform([], seed=seed_stream()))<EOL>log_slice_sample += _log_joint(current_target_log_prob,<EOL>current_momentum)<EOL>reverse_state = current_state<EOL>reverse_target_log_prob = current_target_log_prob<EOL>reverse_grads_target_log_prob = current_grads_target_log_prob<EOL>reverse_momentum = current_momentum<EOL>forward_state = current_state<EOL>forward_target_log_prob = current_target_log_prob<EOL>forward_grads_target_log_prob = current_grads_target_log_prob<EOL>forward_momentum = current_momentum<EOL>next_state = current_state<EOL>next_target_log_prob = current_target_log_prob<EOL>next_grads_target_log_prob = current_grads_target_log_prob<EOL>depth = <NUM_LIT:0><EOL>num_states = <NUM_LIT:1><EOL>continue_trajectory = True<EOL><DEDENT>while continue_trajectory:<EOL><INDENT>direction = tfp.math.random_rademacher([], seed=seed_stream())<EOL>if direction < <NUM_LIT:0>:<EOL><INDENT>[<EOL>reverse_state,<EOL>reverse_target_log_prob,<EOL>reverse_grads_target_log_prob,<EOL>reverse_momentum,<EOL>_,<EOL>_,<EOL>_,<EOL>_,<EOL>next_state_in_subtree,<EOL>next_target_log_prob_in_subtree,<EOL>next_grads_target_log_prob_in_subtree,<EOL>num_states_in_subtree,<EOL>continue_trajectory,<EOL>] = _build_tree(<EOL>value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=reverse_state,<EOL>current_target_log_prob=reverse_target_log_prob,<EOL>current_grads_target_log_prob=reverse_grads_target_log_prob,<EOL>current_momentum=reverse_momentum,<EOL>direction=direction,<EOL>depth=depth,<EOL>step_size=step_size,<EOL>log_slice_sample=log_slice_sample,<EOL>seed=seed_stream())<EOL><DEDENT>else:<EOL><INDENT>[<EOL>_,<EOL>_,<EOL>_,<EOL>_,<EOL>forward_state,<EOL>forward_target_log_prob,<EOL>forward_grads_target_log_prob,<EOL>forward_momentum,<EOL>next_state_in_subtree,<EOL>next_target_log_prob_in_subtree,<EOL>next_grads_target_log_prob_in_subtree,<EOL>num_states_in_subtree,<EOL>continue_trajectory,<EOL>] = _build_tree(<EOL>value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=forward_state,<EOL>current_target_log_prob=forward_target_log_prob,<EOL>current_grads_target_log_prob=forward_grads_target_log_prob,<EOL>current_momentum=forward_momentum,<EOL>direction=direction,<EOL>depth=depth,<EOL>step_size=step_size,<EOL>log_slice_sample=log_slice_sample,<EOL>seed=seed_stream())<EOL><DEDENT>if continue_trajectory:<EOL><INDENT>accept_state_in_subtree = _random_bernoulli(<EOL>[],<EOL>probs=tf.minimum(<NUM_LIT:1.>, num_states_in_subtree / num_states),<EOL>dtype=tf.bool,<EOL>seed=seed_stream())<EOL>if accept_state_in_subtree:<EOL><INDENT>next_state = next_state_in_subtree<EOL>next_target_log_prob = next_target_log_prob_in_subtree<EOL>next_grads_target_log_prob = next_grads_target_log_prob_in_subtree<EOL><DEDENT><DEDENT>has_no_u_turn = tf.logical_and(<EOL>_has_no_u_turn(forward_state, reverse_state, forward_momentum),<EOL>_has_no_u_turn(forward_state, reverse_state, reverse_momentum))<EOL>continue_trajectory = continue_trajectory and has_no_u_turn<EOL>num_states += num_states_in_subtree<EOL>depth += <NUM_LIT:1><EOL><DEDENT>return next_state, next_target_log_prob, next_grads_target_log_prob<EOL><DEDENT>

Simulates a No-U-Turn Sampler (NUTS) trajectory.

    Args:
      target_log_prob_fn: Python callable which takes an argument like
        `*current_state` and returns its (possibly unnormalized) log-density under
        the target distribution.
      current_state: List of `Tensor`s representing the states to simulate from.
      step_size: List of `Tensor`s representing the step sizes for the leapfrog
        integrator. Must have same shape as `current_state`.
      seed: Integer to seed the random number generator.
      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`.
      name: A name for the operation.

    Returns:
      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`.

    Raises:
      NotImplementedError: If the execution mode is not eager.


if depth == <NUM_LIT:0>:  <EOL><INDENT>[<EOL>next_state,<EOL>next_target_log_prob,<EOL>next_grads_target_log_prob,<EOL>next_momentum,<EOL>] = _leapfrog(<EOL>value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=current_state,<EOL>current_grads_target_log_prob=current_grads_target_log_prob,<EOL>current_momentum=current_momentum,<EOL>step_size=direction * step_size)<EOL>next_log_joint = _log_joint(next_target_log_prob, next_momentum)<EOL>num_states = tf.cast(next_log_joint > log_slice_sample, dtype=tf.int32)<EOL>continue_trajectory = (next_log_joint ><EOL>log_slice_sample - max_simulation_error)<EOL>return [<EOL>next_state,<EOL>next_target_log_prob,<EOL>next_grads_target_log_prob,<EOL>next_momentum,<EOL>next_state,<EOL>next_target_log_prob,<EOL>next_grads_target_log_prob,<EOL>next_momentum,<EOL>next_state,<EOL>next_target_log_prob,<EOL>next_grads_target_log_prob,<EOL>num_states,<EOL>continue_trajectory,<EOL>]<EOL><DEDENT>seed_stream = tfd.SeedStream(seed, "<STR_LIT>")<EOL>[<EOL>reverse_state,<EOL>reverse_target_log_prob,<EOL>reverse_grads_target_log_prob,<EOL>reverse_momentum,<EOL>forward_state,<EOL>forward_target_log_prob,<EOL>forward_grads_target_log_prob,<EOL>forward_momentum,<EOL>next_state,<EOL>next_target_log_prob,<EOL>next_grads_target_log_prob,<EOL>num_states,<EOL>continue_trajectory,<EOL>] = _build_tree(value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=current_state,<EOL>current_target_log_prob=current_target_log_prob,<EOL>current_grads_target_log_prob=current_grads_target_log_prob,<EOL>current_momentum=current_momentum,<EOL>direction=direction,<EOL>depth=depth - <NUM_LIT:1>,<EOL>step_size=step_size,<EOL>log_slice_sample=log_slice_sample,<EOL>seed=seed_stream())<EOL>if continue_trajectory:<EOL><INDENT>if direction < <NUM_LIT:0>:<EOL><INDENT>[<EOL>reverse_state,<EOL>reverse_target_log_prob,<EOL>reverse_grads_target_log_prob,<EOL>reverse_momentum,<EOL>_,<EOL>_,<EOL>_,<EOL>_,<EOL>far_state,<EOL>far_target_log_prob,<EOL>far_grads_target_log_prob,<EOL>far_num_states,<EOL>far_continue_trajectory,<EOL>] = _build_tree(<EOL>value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=reverse_state,<EOL>current_target_log_prob=reverse_target_log_prob,<EOL>current_grads_target_log_prob=reverse_grads_target_log_prob,<EOL>current_momentum=reverse_momentum,<EOL>direction=direction,<EOL>depth=depth - <NUM_LIT:1>,<EOL>step_size=step_size,<EOL>log_slice_sample=log_slice_sample,<EOL>seed=seed_stream())<EOL><DEDENT>else:<EOL><INDENT>[<EOL>_,<EOL>_,<EOL>_,<EOL>_,<EOL>forward_state,<EOL>forward_target_log_prob,<EOL>forward_grads_target_log_prob,<EOL>forward_momentum,<EOL>far_state,<EOL>far_target_log_prob,<EOL>far_grads_target_log_prob,<EOL>far_num_states,<EOL>far_continue_trajectory,<EOL>] = _build_tree(<EOL>value_and_gradients_fn=value_and_gradients_fn,<EOL>current_state=forward_state,<EOL>current_target_log_prob=forward_target_log_prob,<EOL>current_grads_target_log_prob=forward_grads_target_log_prob,<EOL>current_momentum=forward_momentum,<EOL>direction=direction,<EOL>depth=depth - <NUM_LIT:1>,<EOL>step_size=step_size,<EOL>log_slice_sample=log_slice_sample,<EOL>seed=seed_stream())<EOL><DEDENT>num_states += far_num_states<EOL>accept_far_state = _random_bernoulli(<EOL>[],<EOL>probs=far_num_states / num_states,<EOL>dtype=tf.bool,<EOL>seed=seed_stream())<EOL>if accept_far_state:<EOL><INDENT>next_state = far_state<EOL>next_target_log_prob = far_target_log_prob<EOL>next_grads_target_log_prob = far_grads_target_log_prob<EOL><DEDENT>has_no_u_turn = tf.logical_and(<EOL>_has_no_u_turn(forward_state, reverse_state, forward_momentum),<EOL>_has_no_u_turn(forward_state, reverse_state, reverse_momentum))<EOL>continue_trajectory = far_continue_trajectory and has_no_u_turn<EOL><DEDENT>return [<EOL>reverse_state,<EOL>reverse_target_log_prob,<EOL>reverse_grads_target_log_prob,<EOL>reverse_momentum,<EOL>forward_state,<EOL>forward_target_log_prob,<EOL>forward_grads_target_log_prob,<EOL>forward_momentum,<EOL>next_state,<EOL>next_target_log_prob,<EOL>next_grads_target_log_prob,<EOL>num_states,<EOL>continue_trajectory,<EOL>]<EOL>

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`.


@functools.wraps(value_and_gradients_fn)<EOL>def func_wrapped(*args, **kwargs):<EOL><INDENT>"""<STR_LIT>"""<EOL>value, grads = value_and_gradients_fn(*args, **kwargs)<EOL>if any(grad is None for grad in grads):<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>return value, grads<EOL><DEDENT>return func_wrapped<EOL>

Wraps value and gradients function to assist with None gradients.


def _has_no_u_turn(state_one, state_two, momentum):

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

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


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


momentum_log_prob = -sum(<EOL>[tf.reduce_sum(input_tensor=<NUM_LIT:0.5> * (m**<NUM_LIT>)) for m in current_momentum])<EOL>return current_target_log_prob + momentum_log_prob<EOL>

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


def _random_bernoulli(shape, probs, dtype=tf.int32, seed=None, name=None):

with tf.compat.v1.name_scope(name, "<STR_LIT>", [shape, probs]):<EOL><INDENT>probs = tf.convert_to_tensor(value=probs)<EOL>random_uniform = tf.random.uniform(shape, dtype=probs.dtype, seed=seed)<EOL>return tf.cast(tf.less(random_uniform, probs), dtype)<EOL><DEDENT>

Returns samples from a Bernoulli distribution.


coeffs = ed.MultivariateNormalDiag(<EOL>loc=tf.zeros(features.shape[<NUM_LIT:1>]), name="<STR_LIT>")<EOL>labels = ed.Bernoulli(<EOL>logits=tf.tensordot(features, coeffs, [[<NUM_LIT:1>], [<NUM_LIT:0>]]), name="<STR_LIT>")<EOL>return labels<EOL>

Bayesian logistic regression, which returns labels given features.


import sklearn.datasets  <EOL>data = sklearn.datasets.covtype.fetch_covtype()<EOL>features = data.data<EOL>labels = data.target<EOL>features -= features.mean(<NUM_LIT:0>)<EOL>features /= features.std(<NUM_LIT:0>)<EOL>features = np.hstack([features, np.ones([features.shape[<NUM_LIT:0>], <NUM_LIT:1>])])<EOL>features = tf.cast(features, dtype=tf.float32)<EOL>_, counts = np.unique(labels, return_counts=True)<EOL>specific_category = np.argmax(counts)<EOL>labels = (labels == specific_category)<EOL>labels = tf.cast(labels, dtype=tf.int32)<EOL>return features, labels<EOL>


def plot_with_expectation(samples, dist, suffix):

_, bins, _ = plt.hist(samples, bins=<NUM_LIT:30>, label="<STR_LIT>")<EOL>bin_width = bins[<NUM_LIT:1>] - bins[<NUM_LIT:0>]<EOL>xs = np.linspace(bins[<NUM_LIT:0>], bins[-<NUM_LIT:1>], <NUM_LIT>)<EOL>pdfs = dist.pdf(xs) * bin_width * len(samples)<EOL>plt.plot(xs, pdfs, label="<STR_LIT>")<EOL>plt.legend()<EOL>ks_stat, pval = scipy.stats.kstest(sorted(samples), dist.cdf)<EOL>plt.title("<STR_LIT>".format(ks_stat, pval))<EOL>savefig(suffix)<EOL>plt.close()<EOL>

Comparison histogram of samples and a line for the expected density.


filename = os.path.join(FLAGS.model_dir, suffix)<EOL>plt.savefig(filename)<EOL>print("<STR_LIT>", filename)<EOL>


def trace(state: State, fn: TransitionOperator, num_steps: IntTensor,<EOL>trace_fn: Callable[[State, TensorNest], TensorNest]<EOL>) -> Tuple[State, TensorNest]:

def fn_wrapper(args, _):<EOL><INDENT>return tf.nest.map_structure(tf.convert_to_tensor, call_fn(fn, args[<NUM_LIT:0>]))<EOL><DEDENT>def trace_fn_wrapper(args):<EOL><INDENT>return tf.nest.map_structure(tf.convert_to_tensor, call_fn(trace_fn, args))<EOL><DEDENT>state = call_fn(fn, state)<EOL>first_trace = trace_fn_wrapper(state)<EOL>state, full_trace = mcmc_util.trace_scan(<EOL>fn_wrapper, state, tf.ones(num_steps - <NUM_LIT:1>), trace_fn=trace_fn_wrapper)<EOL>prepend = lambda x, y: tf.concat(  <EOL>[tf.convert_to_tensor(value=x)[tf.newaxis], y], <NUM_LIT:0>)<EOL>return state, tf.nest.map_structure(prepend, first_trace, full_trace)<EOL>

`TransitionOperator` that runs `fn` repeatedly and traces its outputs.

    Args:
      state: A nest of `Tensor`s or None.
      fn: A `TransitionOperator`.
      num_steps: Number of steps to run the function for. Must be greater than 1.
      trace_fn: Callable that the unpacked outputs of `fn` and returns a nest of
        `Tensor`s. These will be stacked and returned.

    Returns:
      state: The final state returned by `fn`.
      traces: Stacked outputs of `trace_fn`.


def call_fn(fn: TransitionOperator, args: Union[Tuple[Any], Any]) -> Any:

if isinstance(args, (list, tuple)) and not mcmc_util.is_namedtuple_like(args):<EOL><INDENT>args = args  <EOL>return fn(*args)<EOL><DEDENT>else:<EOL><INDENT>return fn(args)<EOL><DEDENT>

Calls a transition operator with args, unpacking args if its a sequence.

    Args:
      fn: A `TransitionOperator`.
      args: Arguments to `fn`

    Returns:
      ret: Return value of `fn`.


def call_and_grads(fn: TransitionOperator, args: Union[Tuple[Any], Any]<EOL>) -> Tuple[tf.Tensor, TensorNest, TensorNest]:

with tf.GradientTape() as tape:<EOL><INDENT>tape.watch(args)<EOL>ret, extra = call_fn(fn, args)<EOL><DEDENT>grads = tape.gradient(ret, args)<EOL>return ret, extra, grads<EOL>

Calls `fn` and returns the gradients with respect to `fn`'s first output.

    Args:
      fn: A `TransitionOperator`.
      args: Arguments to `fn`

    Returns:
      ret: First output of `fn`.
      extra: Second output of `fn`.
      grads: Gradients of `ret` with respect to `args`.


def maybe_broadcast_structure(from_structure: Any, to_structure: Any) -> Any:

flat_from = tf.nest.flatten(from_structure)<EOL>flat_to = tf.nest.flatten(to_structure)<EOL>if len(flat_from) == <NUM_LIT:1>:<EOL><INDENT>flat_from *= len(flat_to)<EOL><DEDENT>return tf.nest.pack_sequence_as(to_structure, flat_from)<EOL>

Maybe broadcasts `from_structure` to `to_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:
      from_structure: A structure.
      to_structure: A structure.

    Returns:
      new_from_structure: Same structure as `to_structure`.


def transform_log_prob_fn(log_prob_fn: PotentialFn,<EOL>bijector: BijectorNest,<EOL>init_state: State = None<EOL>) -> Union[PotentialFn, Tuple[PotentialFn, State]]:

def wrapper(*args):<EOL><INDENT>"""<STR_LIT>"""<EOL>bijector_ = bijector<EOL>args = tf.nest.map_structure(lambda x: <NUM_LIT:0.> + x, args)<EOL>if len(args) == <NUM_LIT:1>:<EOL><INDENT>args = args[<NUM_LIT:0>]<EOL><DEDENT>elif isinstance(bijector_, list):<EOL><INDENT>bijector_ = tuple(bijector_)<EOL><DEDENT>original_space_args = tf.nest.map_structure(lambda b, x: b.forward(x),<EOL>bijector_, args)<EOL>original_space_args = original_space_args  <EOL>original_space_log_prob, extra = call_fn(log_prob_fn, original_space_args)<EOL>event_ndims = tf.nest.map_structure(<EOL>lambda x: tf.rank(x) - tf.rank(original_space_log_prob), args)<EOL>return original_space_log_prob + sum(<EOL>tf.nest.flatten(<EOL>tf.nest.map_structure(<EOL>lambda b, x, e: b.forward_log_det_jacobian(x, event_ndims=e),<EOL>bijector_, args, event_ndims))), [original_space_args, extra]<EOL><DEDENT>if init_state is None:<EOL><INDENT>return wrapper<EOL><DEDENT>else:<EOL><INDENT>return wrapper, tf.nest.map_structure(lambda b, s: b.inverse(s), bijector,<EOL>init_state)<EOL><DEDENT>

Transforms a log-prob function using a bijector.

    This takes a log-prob function and creates a new log-prob function that now
    takes takes state in the domain of the bijector, forward transforms that state
    and calls the original log-prob function. It then returns the log-probability
    that correctly accounts for this transformation.

    The forward-transformed state is pre-pended to the original log-prob
    function's extra returns and returned as the new extra return.

    For convenience you can also pass the initial state (in the original space),
    and this function will return the inverse transformed as the 2nd return value.
    You'd use this to initialize MCMC operators that operate in the transformed
    space.

    Args:
      log_prob_fn: Log prob fn.
      bijector: Bijector(s), must be of the same structure as the `log_prob_fn`
        inputs.
      init_state: Initial state, in the original space.

    Returns:
      transformed_log_prob_fn: Transformed log prob fn.
      transformed_init_state: If `init_state` is provided. Initial state in the
        transformed space.


def leapfrog_step(leapfrog_step_state: LeapFrogStepState,<EOL>step_size: FloatTensor, target_log_prob_fn: PotentialFn,<EOL>kinetic_energy_fn: PotentialFn<EOL>) -> Tuple[LeapFrogStepState, LeapFrogStepExtras]:

state = leapfrog_step_state.state<EOL>state_grads = leapfrog_step_state.state_grads<EOL>momentum = leapfrog_step_state.momentum<EOL>step_size = maybe_broadcast_structure(step_size, state)<EOL>state = tf.nest.map_structure(tf.convert_to_tensor, state)<EOL>momentum = tf.nest.map_structure(tf.convert_to_tensor, momentum)<EOL>state = tf.nest.map_structure(tf.convert_to_tensor, state)<EOL>if state_grads is None:<EOL><INDENT>_, _, state_grads = call_and_grads(target_log_prob_fn, state)<EOL><DEDENT>else:<EOL><INDENT>state_grads = tf.nest.map_structure(tf.convert_to_tensor, state_grads)<EOL><DEDENT>momentum = tf.nest.map_structure(lambda m, sg, s: m + <NUM_LIT:0.5> * sg * s, momentum,<EOL>state_grads, step_size)<EOL>kinetic_energy, kinetic_energy_extra, momentum_grads = call_and_grads(<EOL>kinetic_energy_fn, momentum)<EOL>state = tf.nest.map_structure(lambda x, mg, s: x + mg * s, state,<EOL>momentum_grads, step_size)<EOL>target_log_prob, state_extra, state_grads = call_and_grads(<EOL>target_log_prob_fn, state)<EOL>momentum = tf.nest.map_structure(lambda m, sg, s: m + <NUM_LIT:0.5> * sg * s, momentum,<EOL>state_grads, step_size)<EOL>return LeapFrogStepState(state, state_grads, momentum), LeapFrogStepExtras(<EOL>target_log_prob, state_extra, kinetic_energy, kinetic_energy_extra)<EOL>

Leapfrog `TransitionOperator`.

    Args:
      leapfrog_step_state: LeapFrogStepState.
      step_size: Step size, structure broadcastable to the `target_log_prob_fn`
        state.
      target_log_prob_fn: Target log prob fn.
      kinetic_energy_fn: Kinetic energy fn.

    Returns:
      leapfrog_step_state: LeapFrogStepState.
      leapfrog_step_extras: LeapFrogStepExtras.


def metropolis_hastings_step(current_state: State,<EOL>proposed_state: State,<EOL>energy_change: FloatTensor,<EOL>seed=None) -> Tuple[State, tf.Tensor, tf.Tensor]:

flat_current = tf.nest.flatten(current_state)<EOL>flat_proposed = nest.flatten_up_to(current_state, proposed_state)<EOL>flat_current = [<EOL>p if c is None else c for p, c in zip(flat_proposed, flat_current)<EOL>]<EOL>current_state = tf.nest.pack_sequence_as(current_state, flat_current)<EOL>current_state = tf.nest.map_structure(tf.convert_to_tensor, current_state)<EOL>proposed_state = tf.nest.map_structure(tf.convert_to_tensor, proposed_state)<EOL>energy_change = tf.convert_to_tensor(value=energy_change)<EOL>log_accept_ratio = -energy_change<EOL>log_uniform = tf.math.log(<EOL>tf.random.uniform(<EOL>shape=tf.shape(input=log_accept_ratio),<EOL>dtype=log_accept_ratio.dtype.base_dtype,<EOL>seed=seed))<EOL>is_accepted = log_uniform < log_accept_ratio<EOL>next_state = mcmc_util.choose(<EOL>is_accepted, proposed_state, current_state, name='<STR_LIT>')<EOL>return next_state, is_accepted, log_uniform<EOL>

Metropolis-Hastings step.

    This probabilistically chooses between `current_state` and `proposed_state`
    based on the `energy_change` so as to preserve detailed balance.

    Energy change is the negative of `log_accept_ratio`.

    Args:
      current_state: Current state.
      proposed_state: Proposed state.
      energy_change: E(proposed_state) - E(previous_state).
      seed: For reproducibility.

    Returns:
      new_state: The chosen state.
      is_accepted: Whether the proposed state was accepted.
      log_uniform: The random number that was used to select between the two
        states.


def hamiltonian_monte_carlo(<EOL>hmc_state: HamiltonianMonteCarloState,<EOL>target_log_prob_fn: PotentialFn,<EOL>step_size: Any,<EOL>num_leapfrog_steps: IntTensor,<EOL>momentum: State = None,<EOL>kinetic_energy_fn: PotentialFn = None,<EOL>momentum_sample_fn: MomentumSampleFn = None,<EOL>leapfrog_trace_fn: Callable[[LeapFrogStepState, LeapFrogStepExtras],<EOL>TensorNest] = lambda *args: (),<EOL>seed=None,<EOL>) -> Tuple[HamiltonianMonteCarloState, HamiltonianMonteCarloExtra]:

state = hmc_state.state<EOL>state_grads = hmc_state.state_grads<EOL>target_log_prob = hmc_state.target_log_prob<EOL>state_extra = hmc_state.state_extra<EOL>if kinetic_energy_fn is None:<EOL><INDENT>def kinetic_energy_fn(*momentum):<EOL><INDENT>return tf.add_n([<EOL>tf.reduce_sum(input_tensor=tf.square(x), axis=-<NUM_LIT:1>) / <NUM_LIT><EOL>for x in tf.nest.flatten(momentum)<EOL>]), ()<EOL><DEDENT><DEDENT>if momentum_sample_fn is None:<EOL><INDENT>def momentum_sample_fn(*momentum):<EOL><INDENT>ret = tf.nest.map_structure(<EOL>lambda x: tf.random.normal(tf.shape(input=x), dtype=x.dtype),<EOL>momentum)<EOL>if len(ret) == <NUM_LIT:1>:<EOL><INDENT>return ret[<NUM_LIT:0>]<EOL><DEDENT>else:<EOL><INDENT>return ret<EOL><DEDENT><DEDENT><DEDENT>if momentum is None:<EOL><INDENT>momentum = call_fn(momentum_sample_fn,<EOL>tf.nest.map_structure(tf.zeros_like, state))<EOL><DEDENT>if target_log_prob is None:<EOL><INDENT>target_log_prob, state_extra, state_grads = call_and_grads(<EOL>target_log_prob_fn, state)<EOL><DEDENT>kinetic_energy, _ = call_fn(kinetic_energy_fn, momentum)<EOL>current_energy = -target_log_prob + kinetic_energy<EOL>current_state = HamiltonianMonteCarloState(<EOL>state=state,<EOL>state_grads=state_grads,<EOL>state_extra=state_extra,<EOL>target_log_prob=target_log_prob)<EOL>def leapfrog_wrapper(leapfrog_state, target_log_prob, state_extra):<EOL><INDENT>"""<STR_LIT>"""<EOL>del target_log_prob<EOL>del state_extra<EOL>leapfrog_state, leapfrog_extra = leapfrog_step(<EOL>leapfrog_state,<EOL>step_size=step_size,<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>kinetic_energy_fn=kinetic_energy_fn)<EOL>return [<EOL>leapfrog_state, leapfrog_extra.target_log_prob,<EOL>leapfrog_extra.state_extra<EOL>], leapfrog_extra<EOL><DEDENT>def leapfrog_trace_wrapper_fn(args, leapfrog_extra):<EOL><INDENT>return leapfrog_trace_fn(args[<NUM_LIT:0>], leapfrog_extra)<EOL><DEDENT>leapfrog_wrapper_state = (LeapFrogStepState(state, state_grads, momentum),<EOL>target_log_prob, state_extra)<EOL>[[leapfrog_state, target_log_prob, state_extra], _], leapfrog_trace = trace(<EOL>leapfrog_wrapper_state,<EOL>leapfrog_wrapper,<EOL>num_leapfrog_steps,<EOL>trace_fn=leapfrog_trace_wrapper_fn)<EOL>kinetic_energy, _ = call_fn(kinetic_energy_fn, leapfrog_state.momentum)<EOL>proposed_energy = -target_log_prob + kinetic_energy<EOL>proposed_state = HamiltonianMonteCarloState(<EOL>state=leapfrog_state.state,<EOL>state_grads=leapfrog_state.state_grads,<EOL>target_log_prob=target_log_prob,<EOL>state_extra=state_extra)<EOL>energy_change = proposed_energy - current_energy<EOL>hmc_state, is_accepted, _ = metropolis_hastings_step(<EOL>current_state, proposed_state, energy_change, seed=seed)<EOL>hmc_state = hmc_state  <EOL>return hmc_state, HamiltonianMonteCarloExtra(<EOL>is_accepted=is_accepted,<EOL>proposed_hmc_state=proposed_state,<EOL>log_accept_ratio=-energy_change,<EOL>leapfrog_trace=leapfrog_trace)<EOL>

Hamiltonian Monte Carlo `TransitionOperator`.

    #### Example

    ```python
    step_size = 0.2
    num_steps = 2000
    num_leapfrog_steps = 10
    state = tf.ones([16, 2])

    base_mean = [1., 0]
    base_cov = [[1, 0.5], [0.5, 1]]

    bijector = tfb.Softplus()
    base_dist = tfd.MultivariateNormalFullCovariance(
        loc=base_mean, covariance_matrix=base_cov)
    target_dist = bijector(base_dist)

    def orig_target_log_prob_fn(x):
      return target_dist.log_prob(x), ()

    target_log_prob_fn, state = fun_mcmc.transform_log_prob_fn(
        orig_target_log_prob_fn, bijector, state)

    kernel = tf.function(lambda state: fun_mcmc.hamiltonian_monte_carlo(
        state,
        step_size=step_size,
        num_leapfrog_steps=num_leapfrog_steps,
        target_log_prob_fn=target_log_prob_fn,
        seed=tfp_test_util.test_seed()))

    _, chain = fun_mcmc.trace(
        state=fun_mcmc.HamiltonianMonteCarloState(
            state=state,
            state_grads=None,
            target_log_prob=None,
            state_extra=None),
        fn=kernel,
        num_steps=num_steps,
        trace_fn=lambda state, extra: state.state_extra[0])
    ```

    Args:
      hmc_state: HamiltonianMonteCarloState.
      target_log_prob_fn: Target log prob fn.
      step_size: Step size, structure broadcastable to the `target_log_prob_fn`
        state.
      num_leapfrog_steps: Number of leapfrog steps to take.
      momentum: Initial momentum, passed to `momentum_sample_fn`. Default: zeroes.
      kinetic_energy_fn: Kinetic energy function.
      momentum_sample_fn: Sampler for the momentum.
      leapfrog_trace_fn: Trace function for the leapfrog integrator.
      seed: For reproducibility.

    Returns:
      hmc_state: HamiltonianMonteCarloState
      hmc_extra: HamiltonianMonteCarloExtra


def sign_adaptation(control: FloatNest,<EOL>output: FloatTensor,<EOL>set_point: FloatTensor,<EOL>adaptation_rate: FloatTensor = <NUM_LIT>) -> FloatNest:

def _get_new_control(control, output, set_point):<EOL><INDENT>new_control = mcmc_util.choose(output > set_point,<EOL>control * (<NUM_LIT:1.> + adaptation_rate),<EOL>control / (<NUM_LIT:1.> + adaptation_rate))<EOL>return new_control<EOL><DEDENT>output = maybe_broadcast_structure(output, control)<EOL>set_point = maybe_broadcast_structure(set_point, control)<EOL>return tf.nest.map_structure(_get_new_control, control, output, set_point)<EOL>

A function to do simple sign-based control of a variable.

    ```
    control = control * (1. + adaptation_rate) ** sign(output - set_point)
    ```

    Args:
      control: The control variable.
      output: The output variable.
      set_point: The set point for `output`. This function will adjust `control`
        so that `output` matches `set_point`.
      adaptation_rate: Adaptation rate.

    Returns:
      control: New control.


seed = tfd.SeedStream(seed, '<STR_LIT>')<EOL>def make_iaf():<EOL><INDENT>"""<STR_LIT>"""<EOL>initializer = tf.compat.v2.keras.initializers.VarianceScaling(<EOL><NUM_LIT:2> * <NUM_LIT>, seed=seed() % (<NUM_LIT:2>**<NUM_LIT> - <NUM_LIT:1>))<EOL>made = tfb.AutoregressiveLayer(<EOL>params=<NUM_LIT:2>,<EOL>event_shape=[total_event_size],<EOL>hidden_units=[total_event_size] * num_hidden_layers,<EOL>activation=tf.nn.elu,<EOL>kernel_initializer=initializer,<EOL>dtype=dtype)<EOL>def shift_and_scale(x):<EOL><INDENT>x.set_shape(<EOL>x.shape.merge_with([None] * (x.shape.ndims - <NUM_LIT:1>) + [total_event_size]))<EOL>return tf.unstack(made(x), num=<NUM_LIT:2>, axis=-<NUM_LIT:1>)<EOL><DEDENT>return tfb.Invert(tfb.MaskedAutoregressiveFlow(shift_and_scale))<EOL><DEDENT>def make_swap():<EOL><INDENT>"""<STR_LIT>"""<EOL>permutation = list(reversed(range(total_event_size)))<EOL>return tfb.Permute(permutation)<EOL><DEDENT>bijector = make_iaf()<EOL>bijector = make_swap()(bijector)<EOL>bijector = make_iaf()(bijector)<EOL>bijector = make_swap()(bijector)<EOL>bijector = make_iaf()(bijector)<EOL>bijector = make_swap()(bijector)<EOL>return bijector<EOL>

Creates an stacked IAF bijector.

    This bijector operates on vector-valued events.

    Args:
      total_event_size: Number of dimensions to operate over.
      num_hidden_layers: How many hidden layers to use in each IAF.
      seed: Random seed for the initializers.
      dtype: DType for the variables.

    Returns:
      bijector: The created bijector.


self._parameters = dict(<EOL>target_log_prob_fn=target_log_prob_fn,<EOL>state_shape=state_shape,<EOL>num_step_size_adaptation_steps=num_step_size_adaptation_steps,<EOL>unconstraining_bijector=unconstraining_bijector,<EOL>trainable_bijector_fn=trainable_bijector_fn,<EOL>learning_rate=learning_rate,<EOL>train_batch_size=train_batch_size,<EOL>num_train_steps=num_train_steps,<EOL>train_debug_fn=train_debug_fn,<EOL>seed=seed)<EOL>self._state_shape = tf.nest.map_structure(tf.TensorShape, state_shape)<EOL>if unconstraining_bijector is None:<EOL><INDENT>unconstraining_bijector = tf.nest.map_structure(lambda _: tfb.Identity(),<EOL>self.state_shape)<EOL><DEDENT>self._unconstraining_bijector = unconstraining_bijector<EOL>def _make_reshaped_bijector(b, s):<EOL><INDENT>return tfb.Reshape(<EOL>event_shape_in=s, event_shape_out=[s.num_elements()])(b)(<EOL>tfb.Reshape(event_shape_out=b.inverse_event_shape(s)))<EOL><DEDENT>reshaped_bijector = tf.nest.map_structure(_make_reshaped_bijector,<EOL>unconstraining_bijector,<EOL>self.state_shape)<EOL>blockwise_bijector = tfb.Blockwise(<EOL>bijectors=tf.nest.flatten(reshaped_bijector),<EOL>block_sizes=tf.nest.flatten(<EOL>tf.nest.map_structure(lambda s: s.num_elements(),<EOL>self.state_shape)))<EOL>dtype = dtype_util.common_dtype([learning_rate], preferred_dtype=tf.float32)<EOL>self._dtype = dtype<EOL>trainable_bijector = trainable_bijector_fn(<EOL>self._total_event_size, seed=seed, dtype=dtype)<EOL>self._trainable_bijector = trainable_bijector<EOL>self._bijector = blockwise_bijector(trainable_bijector)<EOL>def flattened_target_log_prob(flat_state):<EOL><INDENT>state = self._unflatten_state(flat_state)<EOL>if isinstance(state, (list, tuple)):<EOL><INDENT>return target_log_prob_fn(*state)<EOL><DEDENT>else:<EOL><INDENT>return target_log_prob_fn(state)<EOL><DEDENT><DEDENT>self._flattened_target_log_prob_val = flattened_target_log_prob<EOL>kernel = tfp.mcmc.HamiltonianMonteCarlo(<EOL>target_log_prob_fn=flattened_target_log_prob,<EOL>step_size=<NUM_LIT:1.>,<EOL>num_leapfrog_steps=self._num_leapfrog_steps(<NUM_LIT:1.>),<EOL>seed=seed if not tf.executing_eagerly() else None)<EOL>kernel = tfp.mcmc.TransformedTransitionKernel(kernel, self._bijector)<EOL>kernel = tfp.mcmc.SimpleStepSizeAdaptation(<EOL>kernel,<EOL>num_adaptation_steps=num_step_size_adaptation_steps,<EOL>target_accept_prob=np.array(<NUM_LIT>, dtype.as_numpy_dtype()))<EOL>self._kernel = kernel<EOL>

Creates the kernel.

        Args:
          target_log_prob_fn: Python callable which takes an argument like
            `current_state` (or `*current_state` if it's a list) and returns its
            (possibly unnormalized) log-density under the target distribution.
          state_shape: A Python list, or list of Python lists (or `TensorShape`s)
            that describes the shape of the problem state. Must be the same
            structure as `current_state`. All shapes must be fully defined.
          num_step_size_adaptation_steps: Number of steps to use for step size
            adaptation. See `tfp.mcmc.SimpleStepSizeAdaptation`.
          unconstraining_bijector: A bijector or list of bijectors that go from
            unconstrained space to the support of the corresponding random variable.
            Must be the same structure as `current_state`.
          trainable_bijector_fn: Creates a trainable, vector-valued event bijector.
            Must be a callable with signature: `(total_event_size, seed) ->
              bijector`, where `total_event_size` is the size of the event.
          learning_rate: Base learning rate to use for training the bijector.
            Internally, learning rate decay is used to stabilize learning.
          train_batch_size: Batch size to use for training the bijector.
          num_train_steps: Number of training steps to train the bijector.
          train_debug_fn: A callable with signature `(NeuTra, step, loss)` called
            for every training step. The first argument is this instance, and `step`
            is the current training step.
          seed: A seed for reproducibility.


@tf.function(autograph=False)<EOL><INDENT>def one_step(self, current_state, previous_kernel_results):<DEDENT>

@tfp.mcmc.internal.util.make_innermost_setter<EOL>def set_num_leapfrog_steps(kernel_results, num_leapfrog_steps):<EOL><INDENT>return kernel_results._replace(<EOL>accepted_results=kernel_results.accepted_results._replace(<EOL>num_leapfrog_steps=num_leapfrog_steps))<EOL><DEDENT>step_size = previous_kernel_results.new_step_size<EOL>previous_kernel_results = set_num_leapfrog_steps(<EOL>previous_kernel_results, self._num_leapfrog_steps(step_size))<EOL>new_state, kernel_results = self._kernel.one_step(<EOL>self._flatten_state(current_state), previous_kernel_results)<EOL>return self._unflatten_state(new_state), kernel_results<EOL>

Runs one iteration of NeuTra.

        Args:
          current_state: `Tensor` or Python `list` of `Tensor`s representing the
            current state(s) of the Markov chain(s). The first `r` dimensions index
            independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
          previous_kernel_results: `collections.namedtuple` containing `Tensor`s
            representing values from previous calls to this function (or from the
            `bootstrap_results` function.)

        Returns:
          next_state: Tensor or Python list of `Tensor`s representing the state(s)
            of the Markov chain(s) after taking exactly one step. Has same type and
            shape as `current_state`.
          kernel_results: `collections.namedtuple` of internal calculations used to
            advance the chain.


def loss():<EOL><INDENT>q = self._flattened_variational_distribution()<EOL>samples = q.sample(self.train_batch_size)<EOL>return tf.reduce_mean(<EOL>input_tensor=q.log_prob(samples) -<EOL>self._flattened_target_log_prob(samples),<EOL>axis=-<NUM_LIT:1>)<EOL><DEDENT>lr = tf.convert_to_tensor(value=self.learning_rate, dtype=self._dtype)<EOL>dtype = lr.dtype<EOL>learning_rate = tf.compat.v2.optimizers.schedules.PiecewiseConstantDecay(<EOL>list(self.num_train_steps *<EOL>np.array([<NUM_LIT>, <NUM_LIT>]).astype(dtype.as_numpy_dtype())),<EOL>[lr, lr * <NUM_LIT:0.1>, lr * <NUM_LIT>])<EOL>opt = tf.compat.v2.optimizers.Adam(learning_rate)<EOL>@tf.function(autograph=False)<EOL>def train_step():<EOL><INDENT>with tf.GradientTape() as tape:<EOL><INDENT>loss_val = loss()<EOL><DEDENT>vals = tape.watched_variables()<EOL>grads = tape.gradient(loss_val, vals)<EOL>grads_and_vals = list(zip(grads, vals))<EOL>opt.apply_gradients(grads_and_vals)<EOL>return loss_val<EOL><DEDENT>for step in range(self.num_train_steps):<EOL><INDENT>loss_val = train_step()<EOL>tf.debugging.assert_all_finite(<EOL>loss_val, '<STR_LIT>'.format(step))<EOL>if self.train_debug_fn:<EOL><INDENT>self.train_debug_fn(self, step, loss_val)<EOL><DEDENT><DEDENT>state_parts = tf.nest.flatten(state)<EOL>flat_state_shapes = tf.nest.flatten(self.state_shape)<EOL>batch_shape = tf.shape(input=state_parts[<NUM_LIT:0>])[:-flat_state_shapes[<NUM_LIT:0>].ndims]<EOL>return self._kernel.bootstrap_results(<EOL>self._flattened_variational_distribution().sample(<EOL>batch_shape, seed=self.seed))<EOL>

Trains the bijector and creates initial `previous_kernel_results`.

        The supplied `state` is only used to determine the number of chains to run
        in parallel_iterations

        Args:
          state: `Tensor` or Python `list` of `Tensor`s representing the initial
            state(s) of the Markov chain(s). The first `r` dimensions index
            independent chains, `r = tf.rank(target_log_prob_fn(*state))`.

        Returns:
          kernel_results: Instance of
            `UncalibratedHamiltonianMonteCarloKernelResults` inside
            `MetropolisHastingsResults` inside `TransformedTransitionKernelResults`
            inside `SimpleStepSizeAdaptationResults`.


def __init__(self, permutation, axis=-<NUM_LIT:1>, validate_args=False, name=None):

with tf.name_scope(name or "<STR_LIT>"):<EOL><INDENT>axis = tf.convert_to_tensor(value=axis, name="<STR_LIT>")<EOL>if not dtype_util.is_integer(axis.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(axis.dtype)))<EOL><DEDENT>permutation = tf.convert_to_tensor(value=permutation, name="<STR_LIT>")<EOL>if not dtype_util.is_integer(permutation.dtype):<EOL><INDENT>raise TypeError("<STR_LIT>".format(<EOL>dtype_util.name(permutation.dtype)))<EOL><DEDENT>p = tf.get_static_value(permutation)<EOL>if p is not None:<EOL><INDENT>if set(p) != set(np.arange(p.size)):<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>p, _ = tf.nn.top_k(<EOL>-permutation, k=tf.shape(input=permutation)[-<NUM_LIT:1>], sorted=True)<EOL>permutation = distribution_util.with_dependencies([<EOL>assert_util.assert_equal(<EOL>-p,<EOL>tf.range(tf.size(input=p)),<EOL>message=("<STR_LIT>"<EOL>"<STR_LIT>")),<EOL>], permutation)<EOL><DEDENT>axis_ = tf.get_static_value(axis)<EOL>if axis_ is None:<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>elif axis_ >= <NUM_LIT:0>:<EOL><INDENT>raise NotImplementedError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>else:<EOL><INDENT>forward_min_event_ndims = int(np.abs(axis_))<EOL><DEDENT>self._permutation = permutation<EOL>self._axis = axis<EOL>super(Permute, self).__init__(<EOL>forward_min_event_ndims=forward_min_event_ndims,<EOL>is_constant_jacobian=True,<EOL>validate_args=validate_args,<EOL>name=name or "<STR_LIT>")<EOL><DEDENT>

Creates the `Permute` bijector.

        Args:
          permutation: An `int`-like vector-shaped `Tensor` representing the
            permutation to apply to the `axis` dimension of the transformed
            `Tensor`.
          axis: Scalar `int` `Tensor` representing the dimension over which to
            `tf.gather`. `axis` must be relative to the end (reading left to right)
            thus must be negative.
            Default value: `-1` (i.e., right-most).
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str`, name given to ops managed by this object.

        Raises:
          TypeError: if `not dtype_util.is_integer(permutation.dtype)`.
          ValueError: if `permutation` does not contain exactly one of each of
            `{0, 1, ..., d}`.
          NotImplementedError: if `axis` is not known prior to graph execution.
          NotImplementedError: if `axis` is not negative.


def _maybe_validate_rightmost_transposed_ndims(<EOL>rightmost_transposed_ndims, validate_args, name=None):

with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>assertions = []<EOL>if not dtype_util.is_integer(rightmost_transposed_ndims.dtype):<EOL><INDENT>raise TypeError('<STR_LIT>')<EOL><DEDENT>if tensorshape_util.rank(rightmost_transposed_ndims.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(rightmost_transposed_ndims.shape) != <NUM_LIT:0>:<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>'.format(<EOL>tensorshape_util.rank(<EOL>rightmost_transposed_ndims.shape)))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions += [assert_util.assert_rank(rightmost_transposed_ndims, <NUM_LIT:0>)]<EOL><DEDENT>rightmost_transposed_ndims_ = tf.get_static_value(<EOL>rightmost_transposed_ndims)<EOL>msg = '<STR_LIT>'<EOL>if rightmost_transposed_ndims_ is not None:<EOL><INDENT>if rightmost_transposed_ndims_ < <NUM_LIT:0>:<EOL><INDENT>raise ValueError(msg[:-<NUM_LIT:1>] + '<STR_LIT>'.format(<EOL>rightmost_transposed_ndims_))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions += [<EOL>assert_util.assert_non_negative(<EOL>rightmost_transposed_ndims, message=msg)<EOL>]<EOL><DEDENT>return assertions<EOL><DEDENT>

Checks that `rightmost_transposed_ndims` is valid.


def _maybe_validate_perm(perm, validate_args, name=None):

with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>assertions = []<EOL>if not dtype_util.is_integer(perm.dtype):<EOL><INDENT>raise TypeError('<STR_LIT>')<EOL><DEDENT>msg = '<STR_LIT>'<EOL>if tensorshape_util.rank(perm.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(perm.shape) != <NUM_LIT:1>:<EOL><INDENT>raise ValueError(<EOL>msg[:-<NUM_LIT:1>] +<EOL>'<STR_LIT>'.format(tensorshape_util.rank(perm.shape)))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions += [assert_util.assert_rank(perm, <NUM_LIT:1>, message=msg)]<EOL><DEDENT>perm_ = tf.get_static_value(perm)<EOL>msg = '<STR_LIT>'<EOL>if perm_ is not None:<EOL><INDENT>if not np.all(np.arange(np.size(perm_)) == np.sort(perm_)):<EOL><INDENT>raise ValueError(msg[:-<NUM_LIT:1>] + '<STR_LIT>'.format(perm_))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions += [<EOL>assert_util.assert_equal(<EOL>tf.sort(perm), tf.range(tf.size(input=perm)), message=msg)<EOL>]<EOL><DEDENT>return assertions<EOL><DEDENT>


def __init__(self, perm=None, rightmost_transposed_ndims=None,<EOL>validate_args=False, name='<STR_LIT>'):

with tf.name_scope(name):<EOL><INDENT>if (rightmost_transposed_ndims is None) == (perm is None):<EOL><INDENT>raise ValueError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>if rightmost_transposed_ndims is not None:<EOL><INDENT>rightmost_transposed_ndims = tf.convert_to_tensor(<EOL>value=rightmost_transposed_ndims,<EOL>dtype=np.int32,<EOL>name='<STR_LIT>')<EOL>rightmost_transposed_ndims_ = tf.get_static_value(<EOL>rightmost_transposed_ndims)<EOL>assertions = _maybe_validate_rightmost_transposed_ndims(<EOL>rightmost_transposed_ndims, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>rightmost_transposed_ndims = tf.identity(rightmost_transposed_ndims)<EOL><DEDENT><DEDENT>perm_start = (<EOL>distribution_util.prefer_static_value(rightmost_transposed_ndims) -<EOL><NUM_LIT:1>)<EOL>perm = tf.range(start=perm_start, limit=-<NUM_LIT:1>, delta=-<NUM_LIT:1>, name='<STR_LIT>')<EOL><DEDENT>else:  <EOL><INDENT>perm = tf.convert_to_tensor(value=perm, dtype=np.int32, name='<STR_LIT>')<EOL>rightmost_transposed_ndims = tf.size(<EOL>input=perm, name='<STR_LIT>')<EOL>rightmost_transposed_ndims_ = tf.get_static_value(<EOL>rightmost_transposed_ndims)<EOL>assertions = _maybe_validate_perm(perm, validate_args)<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>perm = tf.identity(perm)<EOL><DEDENT><DEDENT><DEDENT>if rightmost_transposed_ndims_ is None:<EOL><INDENT>raise NotImplementedError('<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>else:<EOL><INDENT>rightmost_transposed_ndims_ = int(rightmost_transposed_ndims_)<EOL><DEDENT>self._perm = perm<EOL>self._rightmost_transposed_ndims = rightmost_transposed_ndims<EOL>super(Transpose, self).__init__(<EOL>forward_min_event_ndims=rightmost_transposed_ndims_,<EOL>graph_parents=[perm, rightmost_transposed_ndims],<EOL>is_constant_jacobian=True,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Instantiates the `Transpose` bijector.

        Args:
          perm: Positive `int32` vector-shaped `Tensor` representing permutation of
            rightmost dims (for forward transformation).  Note that the `0`th index
            represents the first of the rightmost dims and the largest value must be
            `rightmost_transposed_ndims - 1` and corresponds to `tf.rank(x) - 1`.
            Only one of `perm` and `rightmost_transposed_ndims` can (and must) be
            specified.
            Default value:
            `tf.range(start=rightmost_transposed_ndims, limit=-1, delta=-1)`.
          rightmost_transposed_ndims: Positive `int32` scalar-shaped `Tensor`
            representing the number of rightmost dimensions to permute.
            Only one of `perm` and `rightmost_transposed_ndims` can (and must) be
            specified.
            Default value: `tf.size(perm)`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.

        Raises:
          ValueError: if both or neither `perm` and `rightmost_transposed_ndims` are
            specified.
          NotImplementedError: if `rightmost_transposed_ndims` is not known prior to
            graph execution.


def _event_shape(self, shape, static_perm_to_shape):

rightmost_ = tf.get_static_value(self.rightmost_transposed_ndims)<EOL>if tensorshape_util.rank(shape) is None or rightmost_ is None:<EOL><INDENT>return tf.TensorShape(None)<EOL><DEDENT>if tensorshape_util.rank(shape) < rightmost_:<EOL><INDENT>raise ValueError('<STR_LIT>'.format(<EOL>rightmost_, shape))<EOL><DEDENT>perm_ = tf.get_static_value(self.perm, partial=True)<EOL>if perm_ is None:<EOL><INDENT>return shape[:tensorshape_util.rank(shape) - rightmost_].concatenate(<EOL>[None] * int(rightmost_))<EOL><DEDENT>if sum(p is None for p in perm_) == <NUM_LIT:1>:<EOL><INDENT>present = np.argsort([-<NUM_LIT:1> if p is None else p for p in perm_])<EOL>for i, p in enumerate(present[<NUM_LIT:1>:]):  <EOL><INDENT>if i != p:<EOL><INDENT>perm_ = [i if p is None else p for p in perm_]<EOL>break<EOL><DEDENT><DEDENT><DEDENT>return shape[:tensorshape_util.rank(shape) - rightmost_].concatenate(<EOL>static_perm_to_shape(shape[tensorshape_util.rank(shape) - rightmost_:],<EOL>perm_))<EOL>

Helper for _forward and _inverse_event_shape.


return tf.math.expm1(x)<EOL>

Returns the forward `Bijector` evaluation, i.e., X = g(Y).


return tf.math.log1p(y)<EOL>

Returns the inverse `Bijector` evaluation, i.e., Y = g^-1(X).


return -tf.math.log1p(y)<EOL>

Returns the (log o det o Jacobian o g^-1)(y).


return tf.identity(x)<EOL>

Returns the (log o det o Jacobian o g)(x).


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>graph_parents = []<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>dtype = tf.float32<EOL>if shift is not None:<EOL><INDENT>shift = tf.convert_to_tensor(value=shift, name="<STR_LIT>")<EOL>graph_parents += [shift]<EOL>dtype = dtype_util.base_dtype(shift.dtype)<EOL><DEDENT>self._shift = shift<EOL>if scale is not None:<EOL><INDENT>if (shift is not None and<EOL>not dtype_util.base_equal(shift.dtype, scale.dtype)):<EOL><INDENT>raise TypeError(<EOL>"<STR_LIT>".format(<EOL>shift.dtype, scale.dtype))<EOL><DEDENT>if not isinstance(scale, tf.linalg.LinearOperator):<EOL><INDENT>raise TypeError("<STR_LIT>")<EOL><DEDENT>if validate_args and not scale.is_non_singular:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>graph_parents += scale.graph_parents<EOL>if scale.dtype is not None:<EOL><INDENT>dtype = dtype_util.base_dtype(scale.dtype)<EOL><DEDENT><DEDENT>self._scale = scale<EOL>self._adjoint = adjoint<EOL>super(AffineLinearOperator, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:1>,<EOL>graph_parents=graph_parents,<EOL>is_constant_jacobian=True,<EOL>dtype=dtype,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Instantiates the `AffineLinearOperator` bijector.

        Args:
          shift: Floating-point `Tensor`.
          scale:  Subclass of `LinearOperator`. Represents the (batch) positive
            definite matrix `M` in `R^{k x k}`.
          adjoint: Python `bool` indicating whether to use the `scale` matrix as
            specified or its adjoint.
            Default value: `False`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.

        Raises:
          TypeError: if `scale` is not a `LinearOperator`.
          TypeError: if `shift.dtype` does not match `scale.dtype`.
          ValueError: if not `scale.is_non_singular`.


return self._shift<EOL>

The `shift` `Tensor` in `Y = scale @ X + shift`.


return self._scale<EOL>

The `scale` `LinearOperator` in `Y = scale @ X + shift`.


return self._adjoint<EOL>

`bool` indicating `scale` should be used as conjugate transpose.


def __init__(self, validate_args=False, name="<STR_LIT>"):

self._name = name<EOL>super(Square, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `Square` bijector.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


def _replace_event_shape_in_shape_tensor(<EOL>input_shape, event_shape_in, event_shape_out, validate_args):

output_tensorshape, is_validated = _replace_event_shape_in_tensorshape(<EOL>tensorshape_util.constant_value_as_shape(input_shape),<EOL>event_shape_in,<EOL>event_shape_out)<EOL>validation_dependencies = (<EOL>map(tf.identity, (event_shape_in, event_shape_out))<EOL>if validate_args else ())<EOL>if (tensorshape_util.is_fully_defined(output_tensorshape) and<EOL>(is_validated or not validate_args)):<EOL><INDENT>with tf.control_dependencies(validation_dependencies):<EOL><INDENT>output_shape = tf.convert_to_tensor(<EOL>value=output_tensorshape, name='<STR_LIT>', dtype_hint=tf.int32)<EOL><DEDENT>return output_shape, output_tensorshape<EOL><DEDENT>with tf.control_dependencies(validation_dependencies):<EOL><INDENT>event_shape_in_ndims = (<EOL>tf.size(input=event_shape_in)<EOL>if tensorshape_util.num_elements(event_shape_in.shape) is None else<EOL>tensorshape_util.num_elements(event_shape_in.shape))<EOL>input_non_event_shape, input_event_shape = tf.split(<EOL>input_shape, num_or_size_splits=[-<NUM_LIT:1>, event_shape_in_ndims])<EOL><DEDENT>additional_assertions = []<EOL>if is_validated:<EOL><INDENT>pass<EOL><DEDENT>elif validate_args:<EOL><INDENT>mask = event_shape_in >= <NUM_LIT:0><EOL>explicit_input_event_shape = tf.boolean_mask(<EOL>tensor=input_event_shape, mask=mask)<EOL>explicit_event_shape_in = tf.boolean_mask(<EOL>tensor=event_shape_in, mask=mask)<EOL>additional_assertions.append(<EOL>assert_util.assert_equal(<EOL>explicit_input_event_shape,<EOL>explicit_event_shape_in,<EOL>message='<STR_LIT>'))<EOL><DEDENT>with tf.control_dependencies(additional_assertions):<EOL><INDENT>output_shape = tf.concat([input_non_event_shape, event_shape_out], axis=<NUM_LIT:0>,<EOL>name='<STR_LIT>')<EOL><DEDENT>return output_shape, output_tensorshape<EOL>

Replaces the rightmost dims in a `Tensor` representing a shape.

    Args:
      input_shape: a rank-1 `Tensor` of integers
      event_shape_in: the event shape expected to be present in rightmost dims
        of `shape_in`.
      event_shape_out: the event shape with which to replace `event_shape_in` in
        the rightmost dims of `input_shape`.
      validate_args: Python `bool` indicating whether arguments should
        be checked for correctness.

    Returns:
      output_shape: A rank-1 integer `Tensor` with the same contents as
        `input_shape` except for the event dims, which are replaced with
        `event_shape_out`.


def _replace_event_shape_in_tensorshape(<EOL>input_tensorshape, event_shape_in, event_shape_out):

event_shape_in_ndims = tensorshape_util.num_elements(event_shape_in.shape)<EOL>if tensorshape_util.rank(<EOL>input_tensorshape) is None or event_shape_in_ndims is None:<EOL><INDENT>return tf.TensorShape(None), False  <EOL><DEDENT>input_non_event_ndims = tensorshape_util.rank(<EOL>input_tensorshape) - event_shape_in_ndims<EOL>if input_non_event_ndims < <NUM_LIT:0>:<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'.format(<EOL>tensorshape_util.rank(input_tensorshape), event_shape_in_ndims))<EOL><DEDENT>input_non_event_tensorshape = input_tensorshape[:input_non_event_ndims]<EOL>input_event_tensorshape = input_tensorshape[input_non_event_ndims:]<EOL>event_shape_in_ = tf.get_static_value(event_shape_in)<EOL>is_validated = (<EOL>tensorshape_util.is_fully_defined(input_event_tensorshape) and<EOL>event_shape_in_ is not None)<EOL>if is_validated:<EOL><INDENT>input_event_shape_ = np.int32(input_event_tensorshape)<EOL>mask = event_shape_in_ >= <NUM_LIT:0><EOL>explicit_input_event_shape_ = input_event_shape_[mask]<EOL>explicit_event_shape_in_ = event_shape_in_[mask]<EOL>if not all(explicit_input_event_shape_ == explicit_event_shape_in_):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(input_event_shape_, event_shape_in_))<EOL><DEDENT><DEDENT>event_tensorshape_out = tensorshape_util.constant_value_as_shape(<EOL>event_shape_out)<EOL>if tensorshape_util.rank(event_tensorshape_out) is None:<EOL><INDENT>output_tensorshape = tf.TensorShape(None)<EOL><DEDENT>else:<EOL><INDENT>output_tensorshape = tensorshape_util.concatenate(<EOL>input_non_event_tensorshape, event_tensorshape_out)<EOL><DEDENT>return output_tensorshape, is_validated<EOL>

Replaces the event shape dims of a `TensorShape`.

    Args:
      input_tensorshape: a `TensorShape` instance in which to attempt replacing
        event shape.
      event_shape_in: `Tensor` shape representing the event shape expected to
        be present in (rightmost dims of) `tensorshape_in`. Must be compatible
        with the rightmost dims of `tensorshape_in`.
      event_shape_out: `Tensor` shape representing the new event shape, i.e.,
        the replacement of `event_shape_in`,

    Returns:
      output_tensorshape: `TensorShape` with the rightmost `event_shape_in`
        replaced by `event_shape_out`. Might be partially defined, i.e.,
        `TensorShape(None)`.
      is_validated: Python `bool` indicating static validation happened.

    Raises:
      ValueError: if we can determine the event shape portion of
        `tensorshape_in` as well as `event_shape_in` both statically, and they
        are not compatible. "Compatible" here means that they are identical on
        any dims that are not -1 in `event_shape_in`.


if not dtype_util.is_integer(shape.dtype):<EOL><INDENT>raise TypeError('<STR_LIT>'.format(<EOL>shape, dtype_util.name(shape.dtype)))<EOL><DEDENT>assertions = []<EOL>message = '<STR_LIT>'<EOL>if tensorshape_util.rank(shape.shape) is not None:<EOL><INDENT>if tensorshape_util.rank(shape.shape) > <NUM_LIT:1>:<EOL><INDENT>raise ValueError(message.format(shape))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(assert_util.assert_less(<EOL>tf.rank(shape), <NUM_LIT:2>, message=message.format(shape)))<EOL><DEDENT>shape_ = tf.get_static_value(shape)<EOL>message = '<STR_LIT>'<EOL>if shape_ is not None:<EOL><INDENT>if sum(shape_ == -<NUM_LIT:1>) > <NUM_LIT:1>:<EOL><INDENT>raise ValueError(message.format(shape))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(assert_util.assert_less(<EOL>tf.reduce_sum(input_tensor=tf.cast(tf.equal(shape, -<NUM_LIT:1>), tf.int32)),<EOL><NUM_LIT:2>,<EOL>message=message.format(shape)))<EOL><DEDENT>message = '<STR_LIT>'<EOL>if shape_ is not None:<EOL><INDENT>if np.any(shape_ < -<NUM_LIT:1>):<EOL><INDENT>raise ValueError(message.format(shape))<EOL><DEDENT><DEDENT>elif validate_args:<EOL><INDENT>assertions.append(assert_util.assert_greater(<EOL>shape, -<NUM_LIT:2>, message=message.format(shape)))<EOL><DEDENT>return assertions<EOL>

Check that a shape Tensor is int-type and otherwise sane.


def __init__(self, event_shape_out, event_shape_in=(-<NUM_LIT:1>,),<EOL>validate_args=False, name=None):

with tf.name_scope(name or '<STR_LIT>'):<EOL><INDENT>event_shape_out = tf.convert_to_tensor(<EOL>value=event_shape_out, name='<STR_LIT>', dtype_hint=tf.int32)<EOL>event_shape_in = tf.convert_to_tensor(<EOL>value=event_shape_in, name='<STR_LIT>', dtype_hint=tf.int32)<EOL>forward_min_event_ndims_ = tensorshape_util.num_elements(<EOL>event_shape_in.shape)<EOL>if forward_min_event_ndims_ is None:<EOL><INDENT>raise NotImplementedError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>inverse_min_event_ndims_ = tensorshape_util.num_elements(<EOL>event_shape_out.shape)<EOL>if inverse_min_event_ndims_ is None:<EOL><INDENT>raise NotImplementedError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT>assertions = []<EOL>assertions.extend(_maybe_check_valid_shape(<EOL>event_shape_out, validate_args))<EOL>assertions.extend(_maybe_check_valid_shape(<EOL>event_shape_in, validate_args))<EOL>if assertions:<EOL><INDENT>with tf.control_dependencies(assertions):<EOL><INDENT>event_shape_in = tf.identity(<EOL>event_shape_in, name='<STR_LIT>')<EOL>event_shape_out = tf.identity(<EOL>event_shape_out, name='<STR_LIT>')<EOL><DEDENT><DEDENT>self._event_shape_in = event_shape_in<EOL>self._event_shape_out = event_shape_out<EOL>super(Reshape, self).__init__(<EOL>forward_min_event_ndims=forward_min_event_ndims_,<EOL>inverse_min_event_ndims=inverse_min_event_ndims_,<EOL>is_constant_jacobian=True,<EOL>validate_args=validate_args,<EOL>name=name or '<STR_LIT>')<EOL><DEDENT>

Creates a `Reshape` bijector.

        Args:
          event_shape_out: An `int`-like vector-shaped `Tensor`
            representing the event shape of the transformed output.
          event_shape_in: An optional `int`-like vector-shape `Tensor`
            representing the event shape of the input. This is required in
            order to define inverse operations; the default of (-1,)
            assumes a vector-shaped input.
          validate_args: Python `bool` indicating whether arguments should
            be checked for correctness.
          name: Python `str`, name given to ops managed by this object.

        Raises:
          TypeError: if either `event_shape_in` or `event_shape_out` has
            non-integer `dtype`.
          ValueError: if either of `event_shape_in` or `event_shape_out`
           has non-vector shape (`rank > 1`), or if their sizes do not
           match.


def __init__(self, dct_type=<NUM_LIT:2>, validate_args=False, name='<STR_LIT>'):

<EOL>if dct_type not in (<NUM_LIT:2>, <NUM_LIT:3>):<EOL><INDENT>raise NotImplementedError('<STR_LIT>')<EOL><DEDENT>self._dct_type = dct_type<EOL>super(DiscreteCosineTransform, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:1>,<EOL>inverse_min_event_ndims=<NUM_LIT:1>,<EOL>is_constant_jacobian=True,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `DiscreteCosineTransform` bijector.

        Args:
          dct_type: Python `int`, the DCT type performed by the forward
            transformation. Currently, only 2 and 3 are supported.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


x = self._build_tensor(x)<EOL>for axis in range(-self.evaluate(tf.rank(x)), <NUM_LIT:0>):<EOL><INDENT>bijector = tfb.Cumsum(axis=axis)<EOL>self.assertEqual("<STR_LIT>", bijector.name)<EOL>y = tf.cumsum(x, axis=axis)<EOL>self.assertAllClose(y, self.evaluate(bijector.forward(x)))<EOL>self.assertAllClose(x, self.evaluate(bijector.inverse(y)))<EOL><DEDENT>


event_ndims = int(self.evaluate(tf.rank(x)) - <NUM_LIT:1>)<EOL>self.assertGreaterEqual(event_ndims, <NUM_LIT:1>)<EOL>bijector = tfb.Cumsum(axis=-event_ndims)<EOL>fldj = bijector.forward_log_det_jacobian(<EOL>self._build_tensor(x), event_ndims=event_ndims)<EOL>fldj_theoretical = bijector_test_util.get_fldj_theoretical(<EOL>bijector, x, event_ndims=event_ndims)<EOL>fldj_, fldj_theoretical_ = self.evaluate([fldj, fldj_theoretical])<EOL>self.assertAllEqual(np.zeros_like(fldj_), fldj_)<EOL>self.assertAllClose(np.zeros_like(fldj_theoretical_), fldj_theoretical_)<EOL>


def _validate_block_sizes(block_sizes, bijectors, validate_args):

block_sizes_shape = block_sizes.shape<EOL>if tensorshape_util.is_fully_defined(block_sizes_shape):<EOL><INDENT>if (tensorshape_util.rank(block_sizes_shape) != <NUM_LIT:1> or<EOL>(tensorshape_util.num_elements(block_sizes_shape) != len(bijectors))):<EOL><INDENT>raise ValueError(<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>'.format(block_sizes_shape, len(bijectors)))<EOL><DEDENT>return block_sizes<EOL><DEDENT>elif validate_args:<EOL><INDENT>message = ('<STR_LIT>'<EOL>'<STR_LIT>')<EOL>with tf.control_dependencies([<EOL>assert_util.assert_equal(<EOL>tf.size(input=block_sizes), len(bijectors), message=message),<EOL>assert_util.assert_equal(tf.rank(block_sizes), <NUM_LIT:1>)<EOL>]):<EOL><INDENT>return tf.identity(block_sizes)<EOL><DEDENT><DEDENT>else:<EOL><INDENT>return block_sizes<EOL><DEDENT>


super(Blockwise, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:1>,<EOL>validate_args=validate_args,<EOL>name=name or<EOL>'<STR_LIT>' + '<STR_LIT>'.join([b.name for b in bijectors]))<EOL>if not bijectors:<EOL><INDENT>raise ValueError('<STR_LIT>')<EOL><DEDENT>for bijector in bijectors:<EOL><INDENT>if (bijector.forward_min_event_ndims > <NUM_LIT:1> or<EOL>bijector.inverse_min_event_ndims != bijector.forward_min_event_ndims):<EOL><INDENT>raise NotImplementedError('<STR_LIT>'<EOL>'<STR_LIT>'<EOL>'<STR_LIT>')<EOL><DEDENT><DEDENT>self._bijectors = bijectors<EOL>with self._name_scope('<STR_LIT>'):<EOL><INDENT>if block_sizes is None:<EOL><INDENT>block_sizes = tf.ones(len(bijectors), dtype=tf.int32)<EOL><DEDENT>self._block_sizes = tf.convert_to_tensor(<EOL>value=block_sizes, name='<STR_LIT>', dtype_hint=tf.int32)<EOL>self._block_sizes = _validate_block_sizes(self._block_sizes, bijectors,<EOL>validate_args)<EOL><DEDENT>

Creates the bijector.

        Args:
          bijectors: A non-empty list of bijectors.
          block_sizes: A 1-D integer `Tensor` with each element signifying the
            length of the block of the input vector to pass to the corresponding
            bijector. The length of `block_sizes` must be be equal to the length of
            `bijectors`. If left as None, a vector of 1's is used.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str`, name given to ops managed by this object. Default:
            E.g., `Blockwise([Exp(), Softplus()]).name ==
            'blockwise_of_exp_and_softplus'`.

        Raises:
          NotImplementedError: If a bijector with `event_ndims` > 1 or one that
            reshapes events is passed.
          ValueError: If `bijectors` list is empty.
          ValueError: If size of `block_sizes` does not equal to the length of
            bijectors or is not a vector.


self._graph_parents = []<EOL>self._name = name<EOL>super(MatrixInverseTriL, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:2>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `MatrixInverseTriL` bijector.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


def vector_size_to_square_matrix_size(d, validate_args, name=None):

if isinstance(d, (float, int, np.generic, np.ndarray)):<EOL><INDENT>n = (-<NUM_LIT:1> + np.sqrt(<NUM_LIT:1> + <NUM_LIT:8> * d)) / <NUM_LIT><EOL>if float(int(n)) != n:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>return int(n)<EOL><DEDENT>else:<EOL><INDENT>with tf.name_scope(name or "<STR_LIT>") as name:<EOL><INDENT>n = (-<NUM_LIT:1.> + tf.sqrt(<NUM_LIT:1> + <NUM_LIT> * tf.cast(d, dtype=tf.float32))) / <NUM_LIT><EOL>if validate_args:<EOL><INDENT>with tf.control_dependencies([<EOL>assert_util.assert_equal(<EOL>tf.cast(tf.cast(n, dtype=tf.int32), dtype=tf.float32),<EOL>n,<EOL>message="<STR_LIT>")<EOL>]):<EOL><INDENT>n = tf.identity(n)<EOL><DEDENT><DEDENT>return tf.cast(n, d.dtype)<EOL><DEDENT><DEDENT>


self._upper = upper<EOL>super(FillTriangular, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:1>,<EOL>inverse_min_event_ndims=<NUM_LIT:2>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `FillTriangular` bijector.

        Args:
          upper: Python `bool` representing whether output matrix should be upper
            triangular (`True`) or lower triangular (`False`, default).
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


def __init__(self, axis=-<NUM_LIT:1>, validate_args=False, name="<STR_LIT>"):

if not isinstance(axis, int) or axis >= <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>self._axis = axis<EOL>super(Cumsum, self).__init__(<EOL>is_constant_jacobian=True,<EOL>forward_min_event_ndims=-axis,  <EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `Cumsum` bijector.

        Args:
          axis: Negative Python `int` indicating the axis along which to compute the
            cumulative sum. Note that positive (and zero) values are not supported.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.

        Raises:
          ValueError: If `axis` is not a negative `int`.


return self._axis<EOL>

Returns the axis over which this `Bijector` computes the cumsum.


global INSTANTIABLE_BIJECTORS<EOL>if INSTANTIABLE_BIJECTORS is not None:<EOL><INDENT>return INSTANTIABLE_BIJECTORS<EOL><DEDENT>result = {}<EOL>for (bijector_name, bijector_class) in six.iteritems(tfb.__dict__):<EOL><INDENT>if (not inspect.isclass(bijector_class) or<EOL>not issubclass(bijector_class, tfb.Bijector) or<EOL>bijector_name in SPECIAL_BIJECTORS):<EOL><INDENT>continue<EOL><DEDENT>spec = inspect.getargspec(bijector_class.__init__)<EOL>ctor_args = set(spec.args) | set(<EOL>[arg for arg in (spec.varargs, spec.keywords) if arg is not None])<EOL>unsupported_args = set(ctor_args) - set(['<STR_LIT:name>', '<STR_LIT>', '<STR_LIT>'])<EOL>if unsupported_args:<EOL><INDENT>logging.warning('<STR_LIT>',<EOL>bijector_name, unsupported_args)<EOL>continue<EOL><DEDENT>if not bijector_class()._is_injective:  <EOL><INDENT>logging.warning('<STR_LIT>', bijector_name)<EOL>continue<EOL><DEDENT>result[bijector_name] = (bijector_class,)<EOL><DEDENT>result['<STR_LIT>'] = (tfb.Invert,)<EOL>for bijector_name in sorted(result):<EOL><INDENT>logging.warning('<STR_LIT>', bijector_name)<EOL><DEDENT>INSTANTIABLE_BIJECTORS = result<EOL>return INSTANTIABLE_BIJECTORS<EOL>

Identifies bijectors that are trivially instantiable.


Returns a dict of support mappings for each instantiable bijector.


bijector_names = hps.one_of(map(hps.just, instantiable_bijectors().keys()))<EOL>bijector_name = draw(<EOL>bijector_names.filter(lambda b: (  <EOL>b == '<STR_LIT>' or '<STR_LIT>' in bijector_supports()[b].forward)))<EOL>if bijector_name == '<STR_LIT>':<EOL><INDENT>underlying = draw(<EOL>bijector_names.filter(<EOL>lambda b: '<STR_LIT>' in bijector_supports()[b].inverse))<EOL>underlying = instantiable_bijectors()[underlying][<NUM_LIT:0>](validate_args=True)<EOL>return tfb.Invert(underlying, validate_args=True)<EOL><DEDENT>return instantiable_bijectors()[bijector_name][<NUM_LIT:0>](validate_args=True)<EOL>

Draws bijectors which can act on [numerically] unconstrained events.


if isinstance(bijector, tfb.CholeskyToInvCholesky):<EOL><INDENT>def additional_check(dist):<EOL><INDENT>return (tensorshape_util.rank(dist.event_shape) == <NUM_LIT:2> and<EOL>int(dist.event_shape[<NUM_LIT:0>]) == int(dist.event_shape[<NUM_LIT:1>]))<EOL><DEDENT><DEDENT>else:<EOL><INDENT>additional_check = lambda dist: True<EOL><DEDENT>def distribution_filter(dist):<EOL><INDENT>if not dtype_util.is_floating(dist.dtype):<EOL><INDENT>return False<EOL><DEDENT>if bijector.forward_min_event_ndims > tensorshape_util.rank(<EOL>dist.event_shape):<EOL><INDENT>return False<EOL><DEDENT>return additional_check(dist)<EOL><DEDENT>return distribution_filter<EOL>

Returns filter function f s.t. f(dist)=True => bijector can act on dist.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>concentration1 = self._maybe_assert_valid_concentration(<EOL>tf.convert_to_tensor(value=concentration1, name="<STR_LIT>"),<EOL>validate_args=validate_args)<EOL>concentration0 = self._maybe_assert_valid_concentration(<EOL>tf.convert_to_tensor(value=concentration0, name="<STR_LIT>"),<EOL>validate_args=validate_args)<EOL><DEDENT>self._concentration1 = concentration1<EOL>self._concentration0 = concentration0<EOL>super(Kumaraswamy, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `Kumaraswamy` bijector.

        Args:
          concentration1: Python `float` scalar indicating the transform power,
            i.e., `Y = g(X) = (1 - (1 - X)**(1 / b))**(1 / a)` where `a` is
            `concentration1`.
          concentration0: Python `float` scalar indicating the transform power,
            i.e., `Y = g(X) = (1 - (1 - X)**(1 / b))**(1 / a)` where `b` is
            `concentration0`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.


return self._concentration1<EOL>


return self._concentration0<EOL>


def _maybe_assert_valid_concentration(self, concentration, validate_args):

if not validate_args:<EOL><INDENT>return concentration<EOL><DEDENT>return distribution_util.with_dependencies([<EOL>assert_util.assert_positive(<EOL>concentration, message="<STR_LIT>"),<EOL>], concentration)<EOL>

Checks the validity of a concentration parameter.


def __init__(self, bijector, validate_args=False, name=None):

if not bijector._is_injective:  <EOL><INDENT>raise NotImplementedError(<EOL>"<STR_LIT>")<EOL><DEDENT>self._bijector = bijector<EOL>super(Invert, self).__init__(<EOL>graph_parents=bijector.graph_parents,<EOL>forward_min_event_ndims=bijector.inverse_min_event_ndims,<EOL>inverse_min_event_ndims=bijector.forward_min_event_ndims,<EOL>is_constant_jacobian=bijector.is_constant_jacobian,<EOL>validate_args=validate_args,<EOL>dtype=bijector.dtype,<EOL>name=name or "<STR_LIT:_>".join(["<STR_LIT>", bijector.name]))<EOL>

Creates a `Bijector` which swaps the meaning of `inverse` and `forward`.

        Note: An inverted bijector's `inverse_log_det_jacobian` is often more
        efficient if the base bijector implements `_forward_log_det_jacobian`. If
        `_forward_log_det_jacobian` is not implemented then the following code is
        used:

        ```python
        y = self.inverse(x, **kwargs)
        return -self.inverse_log_det_jacobian(y, **kwargs)
        ```

        Args:
          bijector: Bijector instance.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str`, name given to ops managed by this object.


self._graph_parents = []<EOL>self._name = name<EOL>self._validate_args = validate_args<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>power = tf.get_static_value(<EOL>tf.convert_to_tensor(value=power, name="<STR_LIT>"))<EOL><DEDENT>if power is None or power < <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>")<EOL><DEDENT>self._power = power<EOL>super(PowerTransform, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>

Instantiates the `PowerTransform` bijector.

        Args:
          power: Python `float` scalar indicating the transform power, i.e.,
            `Y = g(X) = (1 + X * c)**(1 / c)` where `c` is the `power`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          name: Python `str` name given to ops managed by this object.

        Raises:
          ValueError: if `power < 0` or is not known statically.


return self._power<EOL>


super(Inline, self).__init__(<EOL>forward_min_event_ndims=forward_min_event_ndims,<EOL>inverse_min_event_ndims=inverse_min_event_ndims,<EOL>is_constant_jacobian=is_constant_jacobian,<EOL>validate_args=validate_args,<EOL>name=name)<EOL>self._forward_fn = forward_fn<EOL>self._inverse_fn = inverse_fn<EOL>self._inverse_log_det_jacobian_fn = inverse_log_det_jacobian_fn<EOL>self._forward_log_det_jacobian_fn = forward_log_det_jacobian_fn<EOL>self._forward_event_shape_fn = forward_event_shape_fn<EOL>self._forward_event_shape_tensor_fn = forward_event_shape_tensor_fn<EOL>self._inverse_event_shape_fn = inverse_event_shape_fn<EOL>self._inverse_event_shape_tensor_fn = inverse_event_shape_tensor_fn<EOL>

Creates a `Bijector` from callables.

        Args:
          forward_fn: Python callable implementing the forward transformation.
          inverse_fn: Python callable implementing the inverse transformation.
          inverse_log_det_jacobian_fn: Python callable implementing the
            log o det o jacobian of the inverse transformation.
          forward_log_det_jacobian_fn: Python callable implementing the
            log o det o jacobian of the forward transformation.
          forward_event_shape_fn: Python callable implementing non-identical
            static event shape changes. Default: shape is assumed unchanged.
          forward_event_shape_tensor_fn: Python callable implementing non-identical
            event shape changes. Default: shape is assumed unchanged.
          inverse_event_shape_fn: Python callable implementing non-identical
            static event shape changes. Default: shape is assumed unchanged.
          inverse_event_shape_tensor_fn: Python callable implementing non-identical
            event shape changes. Default: shape is assumed unchanged.
          is_constant_jacobian: Python `bool` indicating that the Jacobian is
            constant for all input arguments.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
          forward_min_event_ndims: Python `int` indicating the minimal
            dimensionality this bijector acts on.
          inverse_min_event_ndims: Python `int` indicating the minimal
            dimensionality this bijector acts on.
          name: Python `str`, name given to ops managed by this object.


with tf.name_scope(name) as name:<EOL><INDENT>if diag_bijector is None:<EOL><INDENT>diag_bijector = softplus.Softplus(validate_args=validate_args)<EOL><DEDENT>if diag_shift is not None:<EOL><INDENT>diag_shift = tf.convert_to_tensor(<EOL>value=diag_shift, dtype=diag_bijector.dtype, name="<STR_LIT>")<EOL>diag_bijector = chain.Chain([<EOL>affine_scalar.AffineScalar(shift=diag_shift),<EOL>diag_bijector<EOL>])<EOL><DEDENT>super(ScaleTriL, self).__init__(<EOL>[transform_diagonal.TransformDiagonal(diag_bijector=diag_bijector),<EOL>fill_triangular.FillTriangular()],<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Instantiates the `ScaleTriL` bijector.

        Args:
          diag_bijector: `Bijector` instance, used to transform the output diagonal
            to be positive.
            Default value: `None` (i.e., `tfb.Softplus()`).
          diag_shift: Float value broadcastable and added to all diagonal entries
            after applying the `diag_bijector`. Setting a positive
            value forces the output diagonal entries to be positive, but
            prevents inverting the transformation for matrices with
            diagonal entries less than this value.
            Default value: `1e-5`.
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness.
            Default value: `False` (i.e., arguments are not validated).
          name: Python `str` name given to ops managed by this object.
            Default value: `scale_tril`.


self._graph_parents = []<EOL>self._name = name<EOL>with self._name_scope("<STR_LIT>"):<EOL><INDENT>super(AbsoluteValue, self).__init__(<EOL>forward_min_event_ndims=<NUM_LIT:0>,<EOL>validate_args=validate_args,<EOL>name=name)<EOL><DEDENT>

Instantiates the `AbsoluteValue` bijector.

        Args:
          validate_args: Python `bool` indicating whether arguments should be
            checked for correctness, in particular whether inputs to `inverse` and
            `inverse_log_det_jacobian` are non-negative.
          name: Python `str` name given to ops managed by this object.


def __new__(cls, x=None, y=None, ildj=None, kwargs=None):

return super(_Mapping, cls).__new__(cls, x, y, ildj, kwargs)<EOL>

Custom __new__ so namedtuple items have defaults.

        Args:
          x: `Tensor` or None. Input to forward; output of inverse.
          y: `Tensor` or None. Input to inverse; output of forward.
          ildj: `Tensor`. This is the (un-reduce_sum'ed) inverse log det jacobian.
          kwargs: Python dictionary. Extra args supplied to forward/inverse/etc
            functions.

        Returns:
          mapping: New instance of _Mapping.


return self._deep_tuple(self.kwargs)<EOL>

Returns subkey used for caching (nested under either `x` or `y`).


def merge(self, x=None, y=None, ildj=None, kwargs=None, mapping=None):

if mapping is None:<EOL><INDENT>mapping = _Mapping(x=x, y=y, ildj=ildj, kwargs=kwargs)<EOL><DEDENT>elif any(arg is not None for arg in [x, y, ildj, kwargs]):<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>return _Mapping(<EOL>x=self._merge(self.x, mapping.x),<EOL>y=self._merge(self.y, mapping.y),<EOL>ildj=self._merge(self.ildj, mapping.ildj),<EOL>kwargs=self._merge(self.kwargs, mapping.kwargs, use_equals=True))<EOL>

Returns new _Mapping with args merged with self.

        Args:
          x: `Tensor` or None. Input to forward; output of inverse.
          y: `Tensor` or None. Input to inverse; output of forward.
          ildj: `Tensor`. This is the (un-reduce_sum'ed) inverse log det jacobian.
          kwargs: Python dictionary. Extra args supplied to forward/inverse/etc
            functions.
          mapping: Instance of _Mapping to merge. Can only be specified if no other
            arg is specified.

        Returns:
          mapping: New instance of `_Mapping` which has inputs merged with self.

        Raises:
          ValueError: if mapping and any other arg is not `None`.


return _Mapping(<EOL>x=None if field == "<STR_LIT:x>" else self.x,<EOL>y=None if field == "<STR_LIT:y>" else self.y,<EOL>ildj=self.ildj,<EOL>kwargs=self.kwargs)<EOL>

To support weak referencing, removes cache key from the cache value.


def _merge(self, old, new, use_equals=False):

if old is None:<EOL><INDENT>return new<EOL><DEDENT>if new is None:<EOL><INDENT>return old<EOL><DEDENT>if (old == new) if use_equals else (old is new):<EOL><INDENT>return old<EOL><DEDENT>raise ValueError("<STR_LIT>" % (old, new))<EOL>

Helper to merge which handles merging one value.


if isinstance(x, dict):<EOL><INDENT>return self._deep_tuple(tuple(sorted(x.items())))<EOL><DEDENT>elif isinstance(x, (list, tuple)):<EOL><INDENT>return tuple(map(self._deep_tuple, x))<EOL><DEDENT>return x<EOL>

Converts nested `tuple`, `list`, or `dict` to nested `tuple`.


self._graph_parents = graph_parents or []<EOL>if forward_min_event_ndims is None and inverse_min_event_ndims is None:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>elif inverse_min_event_ndims is None:<EOL><INDENT>inverse_min_event_ndims = forward_min_event_ndims<EOL><DEDENT>elif forward_min_event_ndims is None:<EOL><INDENT>forward_min_event_ndims = inverse_min_event_ndims<EOL><DEDENT>if not isinstance(forward_min_event_ndims, int):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>type(forward_min_event_ndims).__name__))<EOL><DEDENT>if not isinstance(inverse_min_event_ndims, int):<EOL><INDENT>raise TypeError("<STR_LIT>"<EOL>"<STR_LIT>".format(<EOL>type(inverse_min_event_ndims).__name__))<EOL><DEDENT>if forward_min_event_ndims < <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>if inverse_min_event_ndims < <NUM_LIT:0>:<EOL><INDENT>raise ValueError("<STR_LIT>"<EOL>"<STR_LIT>")<EOL><DEDENT>self._forward_min_event_ndims = forward_min_event_ndims<EOL>self._inverse_min_event_ndims = inverse_min_event_ndims<EOL>self._is_constant_jacobian = is_constant_jacobian<EOL>self._constant_ildj = {}<EOL>self._validate_args = validate_args<EOL>self._dtype = dtype<EOL>self._from_y = WeakKeyDefaultDict()<EOL>self._from_x = WeakKeyDefaultDict()<EOL>if name:<EOL><INDENT>self._name = name<EOL><DEDENT>else:<EOL><INDENT>def camel_to_snake(name):<EOL><INDENT>s1 = re.sub("<STR_LIT>", r"<STR_LIT>", name)<EOL>return re.sub("<STR_LIT>", r"<STR_LIT>", s1).lower()<EOL><DEDENT>self._name = camel_to_snake(type(self).__name__.lstrip("<STR_LIT:_>"))<EOL><DEDENT>for i, t in enumerate(self._graph_parents):<EOL><INDENT>if t is None or not tf.is_tensor(t):<EOL><INDENT>raise ValueError("<STR_LIT>" % (i, t))<EOL><DEDENT><DEDENT>

Constructs Bijector.

        A `Bijector` transforms random variables into new random variables.

        Examples:

        ```python
        # Create the Y = g(X) = X transform.
        identity = Identity()

        # Create the Y = g(X) = exp(X) transform.
        exp = Exp()
        ```

        See `Bijector` subclass docstring for more details and specific examples.

        Args:
          graph_parents: Python list of graph prerequisites of this `Bijector`.
          is_constant_jacobian: Python `bool` indicating that the Jacobian matrix is
            not a function of the input.
          validate_args: Python `bool`, default `False`. Whether to validate input
            with asserts. If `validate_args` is `False`, and the inputs are invalid,
            correct behavior is not guaranteed.
          dtype: `tf.dtype` supported by this `Bijector`. `None` means dtype is not
            enforced.
          forward_min_event_ndims: Python `integer` indicating the minimum number of
            dimensions `forward` operates on.
          inverse_min_event_ndims: Python `integer` indicating the minimum number of
            dimensions `inverse` operates on. Will be set to
            `forward_min_event_ndims` by default, if no value is provided.
          name: The name to give Ops created by the initializer.

        Raises:
          ValueError:  If neither `forward_min_event_ndims` and
            `inverse_min_event_ndims` are specified, or if either of them is
            negative.
          ValueError:  If a member of `graph_parents` is not a `Tensor`.


return self._graph_parents<EOL>

Returns this `Bijector`'s graph_parents as a Python list.


return self._forward_min_event_ndims<EOL>

Returns the minimal number of dimensions bijector.forward operates on.


return self._inverse_min_event_ndims<EOL>

Returns the minimal number of dimensions bijector.inverse operates on.


return self._is_constant_jacobian<EOL>

Returns true iff the Jacobian matrix is not a function of x.

        Note: Jacobian matrix is either constant for both forward and inverse or
        neither.

        Returns:
          is_constant_jacobian: Python `bool`.


return True<EOL>

Returns true iff the forward map `g` is injective (one-to-one function).

        **WARNING** This hidden property and its behavior are subject to change.

        Note:  Non-injective maps `g` are supported, provided their domain `D` can
        be partitioned into `k` disjoint subsets, `Union{D1, ..., Dk}`, such that,
        ignoring sets of measure zero, the restriction of `g` to each subset is a
        differentiable bijection onto `g(D)`.

        Returns:
          is_injective: Python `bool`.


return self._validate_args<EOL>

Returns True if Tensor arguments will be validated.


return self._dtype<EOL>

dtype of `Tensor`s transformable by this distribution.


return self._name<EOL>

Returns the string name of this `Bijector`.


def __call__(self, value, name=None, **kwargs):

<EOL>from tensorflow_probability.python.bijectors import chain  <EOL>from tensorflow_probability.python.distributions import distribution  <EOL>from tensorflow_probability.python.distributions import transformed_distribution  <EOL>if type(value) is transformed_distribution.TransformedDistribution:  <EOL><INDENT>new_kwargs = value.parameters<EOL>new_kwargs.update(kwargs)<EOL>new_kwargs["<STR_LIT:name>"] = name or new_kwargs.get("<STR_LIT:name>", None)<EOL>new_kwargs["<STR_LIT>"] = self(value.bijector)<EOL>return transformed_distribution.TransformedDistribution(**new_kwargs)<EOL><DEDENT>if isinstance(value, distribution.Distribution):<EOL><INDENT>return transformed_distribution.TransformedDistribution(<EOL>distribution=value,<EOL>bijector=self,<EOL>name=name,<EOL>**kwargs)<EOL><DEDENT>if isinstance(value, chain.Chain):<EOL><INDENT>new_kwargs = kwargs.copy()<EOL>new_kwargs["<STR_LIT>"] = [self] + ([] if value.bijectors is None<EOL>else list(value.bijectors))<EOL>if "<STR_LIT>" not in new_kwargs:<EOL><INDENT>new_kwargs["<STR_LIT>"] = value.validate_args<EOL><DEDENT>new_kwargs["<STR_LIT:name>"] = name or value.name<EOL>return chain.Chain(**new_kwargs)<EOL><DEDENT>if isinstance(value, Bijector):<EOL><INDENT>return chain.Chain([self, value], name=name, **kwargs)<EOL><DEDENT>return self._call_forward(value, name=name or "<STR_LIT>", **kwargs)<EOL>

Applies or composes the `Bijector`, depending on input type.

        This is a convenience function which applies the `Bijector` instance in
        three different ways, depending on the input:

        1. If the input is a `tfd.Distribution` instance, return
           `tfd.TransformedDistribution(distribution=input, bijector=self)`.
        2. If the input is a `tfb.Bijector` instance, return
           `tfb.Chain([self, input])`.
        3. Otherwise, return `self.forward(input)`

        Args:
          value: A `tfd.Distribution`, `tfb.Bijector`, or a `Tensor`.
          name: Python `str` name given to ops created by this function.
          **kwargs: Additional keyword arguments passed into the created
            `tfd.TransformedDistribution`, `tfb.Bijector`, or `self.forward`.

        Returns:
          composition: A `tfd.TransformedDistribution` if the input was a
            `tfd.Distribution`, a `tfb.Chain` if the input was a `tfb.Bijector`, or
            a `Tensor` computed by `self.forward`.

        #### Examples

        ```python
        sigmoid = tfb.Reciprocal()(
            tfb.AffineScalar(shift=1.)(
              tfb.Exp()(
                tfb.AffineScalar(scale=-1.))))
        # ==> `tfb.Chain([
        #         tfb.Reciprocal(),
        #         tfb.AffineScalar(shift=1.),
        #         tfb.Exp(),
        #         tfb.AffineScalar(scale=-1.),
        #      ])`  # ie, `tfb.Sigmoid()`

        log_normal = tfb.Exp()(tfd.Normal(0, 1))
        # ==> `tfd.TransformedDistribution(tfd.Normal(0, 1), tfb.Exp())`

        tfb.Exp()([-1., 0., 1.])
        # ==> tf.exp([-1., 0., 1.])
        ```


<EOL>return input_shape<EOL>

Subclass implementation for `forward_event_shape_tensor` function.


with self._name_scope(name):<EOL><INDENT>input_shape = tf.convert_to_tensor(<EOL>value=input_shape, dtype=tf.int32, name="<STR_LIT>")<EOL>return self._forward_event_shape_tensor(input_shape)<EOL><DEDENT>

Shape of a single sample from a single batch as an `int32` 1D `Tensor`.

        Args:
          input_shape: `Tensor`, `int32` vector indicating event-portion shape
            passed into `forward` function.
          name: name to give to the op

        Returns:
          forward_event_shape_tensor: `Tensor`, `int32` vector indicating
            event-portion shape after applying `forward`.


<EOL>return input_shape<EOL>

Subclass implementation for `forward_event_shape` public function.


return self._forward_event_shape(tf.TensorShape(input_shape))<EOL>

Shape of a single sample from a single batch as a `TensorShape`.

        Same meaning as `forward_event_shape_tensor`. May be only partially defined.

        Args:
          input_shape: `TensorShape` indicating event-portion shape passed into
            `forward` function.

        Returns:
          forward_event_shape_tensor: `TensorShape` indicating event-portion shape
            after applying `forward`. Possibly unknown.


<EOL>return output_shape<EOL>

Subclass implementation for `inverse_event_shape_tensor` function.


with self._name_scope(name):<EOL><INDENT>output_shape = tf.convert_to_tensor(<EOL>value=output_shape, dtype=tf.int32, name="<STR_LIT>")<EOL>return self._inverse_event_shape_tensor(output_shape)<EOL><DEDENT>

Shape of a single sample from a single batch as an `int32` 1D `Tensor`.

        Args:
          output_shape: `Tensor`, `int32` vector indicating event-portion shape
            passed into `inverse` function.
          name: name to give to the op

        Returns:
          inverse_event_shape_tensor: `Tensor`, `int32` vector indicating
            event-portion shape after applying `inverse`.


<EOL>return tf.TensorShape(output_shape)<EOL>

Subclass implementation for `inverse_event_shape` public function.


return self._inverse_event_shape(tf.TensorShape(output_shape))<EOL>

Shape of a single sample from a single batch as a `TensorShape`.

        Same meaning as `inverse_event_shape_tensor`. May be only partially defined.

        Args:
          output_shape: `TensorShape` indicating event-portion shape passed into
            `inverse` function.

        Returns:
          inverse_event_shape_tensor: `TensorShape` indicating event-portion shape
            after applying `inverse`. Possibly unknown.


raise NotImplementedError("<STR_LIT>")<EOL>

Subclass implementation for `forward` public function.


def _call_forward(self, x, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>x = tf.convert_to_tensor(value=x, name="<STR_LIT:x>")<EOL>self._maybe_assert_dtype(x)<EOL>if not self._is_injective:  <EOL><INDENT>return self._forward(x, **kwargs)<EOL><DEDENT>mapping = self._lookup(x=x, kwargs=kwargs)<EOL>if mapping.y is not None:<EOL><INDENT>return mapping.y<EOL><DEDENT>mapping = mapping.merge(y=self._forward(x, **kwargs))<EOL>self._cache_by_y(mapping)<EOL>if not tf.executing_eagerly():<EOL><INDENT>self._cache_by_x(mapping)<EOL><DEDENT>return mapping.y<EOL><DEDENT>

Wraps call to _forward, allowing extra shared logic.


def forward(self, x, name="<STR_LIT>", **kwargs):

return self._call_forward(x, name, **kwargs)<EOL>

Returns the forward `Bijector` evaluation, i.e., X = g(Y).

        Args:
          x: `Tensor`. The input to the "forward" evaluation.
          name: The name to give this op.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          `Tensor`.

        Raises:
          TypeError: if `self.dtype` is specified and `x.dtype` is not
            `self.dtype`.
          NotImplementedError: if `_forward` is not implemented.


Subclass implementation for `inverse` public function.


def _call_inverse(self, y, name, **kwargs):

with self._name_scope(name):<EOL><INDENT>y = tf.convert_to_tensor(value=y, name="<STR_LIT:y>")<EOL>self._maybe_assert_dtype(y)<EOL>if not self._is_injective:  <EOL><INDENT>return self._inverse(y, **kwargs)<EOL><DEDENT>mapping = self._lookup(y=y, kwargs=kwargs)<EOL>if mapping.x is not None:<EOL><INDENT>return mapping.x<EOL><DEDENT>mapping = mapping.merge(x=self._inverse(y, **kwargs))<EOL>self._cache_by_x(mapping)<EOL>if not tf.executing_eagerly():<EOL><INDENT>self._cache_by_y(mapping)<EOL><DEDENT>return mapping.x<EOL><DEDENT>

Wraps call to _inverse, allowing extra shared logic.


def inverse(self, y, name="<STR_LIT>", **kwargs):

return self._call_inverse(y, name, **kwargs)<EOL>

Returns the inverse `Bijector` evaluation, i.e., X = g^{-1}(Y).

        Args:
          y: `Tensor`. The input to the "inverse" evaluation.
          name: The name to give this op.
          **kwargs: Named arguments forwarded to subclass implementation.

        Returns:
          `Tensor`, if this bijector is injective.
            If not injective, returns the k-tuple containing the unique
            `k` points `(x1, ..., xk)` such that `g(xi) = y`.

        Raises:
          TypeError: if `self.dtype` is specified and `y.dtype` is not
            `self.dtype`.
          NotImplementedError: if `_inverse` is not implemented.