<commit_before># Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Piecewise Rational Quadratic Spline bijector."""

import collections
import functools

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.bijectors import bijector
from tensorflow_probability.python.bijectors import softplus as softplus_bijector
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import parameter_properties
from tensorflow_probability.python.internal import tensor_util
from tensorflow_probability.python.internal import tensorshape_util


def _ensure_at_least_1d(t):
  t = tf.convert_to_tensor(t)
  return t + tf.zeros([1], dtype=t.dtype)


def _padded(t, lhs, rhs=None):
  """Left pads and optionally right pads the innermost axis of `t`."""
  lhs = tf.convert_to_tensor(lhs, dtype=t.dtype)
  zeros = tf.zeros([tf.rank(t) - 1, 2], dtype=tf.int32)
  lhs_paddings = tf.concat([zeros, [[1, 0]]], axis=0)
  result = tf.pad(t, paddings=lhs_paddings, constant_values=lhs)
  if rhs is not None:
    rhs = tf.convert_to_tensor(rhs, dtype=t.dtype)
    rhs_paddings = tf.concat([zeros, [[0, 1]]], axis=0)
    result = tf.pad(result, paddings=rhs_paddings, constant_values=rhs)
  return result


def _knot_positions(bin_sizes, range_min):
  return _padded(tf.cumsum(bin_sizes, axis=-1) + range_min, lhs=range_min)


_SplineShared = collections.namedtuple(
    'SplineShared', 'out_of_bounds,x_k,y_k,d_k,d_kp1,h_k,w_k,s_k')


class RationalQuadraticSpline(bijector.AutoCompositeTensorBijector):
  """A piecewise rational quadratic spline, as developed in [1].

  This transformation represents a monotonically increasing piecewise rational
  quadratic function. Outside of the bounds of `knot_x`/`knot_y`, the transform
  behaves as an identity function.

  Typically this bijector will be used as part of a chain, with splines for
  trailing `x` dimensions conditioned on some of the earlier `x` dimensions, and
  with the inverse then solved first for unconditioned dimensions, then using
  conditioning derived from those inverses, and so forth. For example, if we
  split a 15-D `xs` vector into 3 components, we may implement a forward and
  inverse as follows:

  ```python
  nsplits = 3

  class SplineParams(tf.Module):

    def __init__(self, nbins=32):
      self._nbins = nbins
      self._built = False
      self._bin_widths = None
      self._bin_heights = None
      self._knot_slopes = None

    def __call__(self, x, nunits):
      if not self._built:
        def _bin_positions(x):
          out_shape = tf.concat((tf.shape(x)[:-1], (nunits, self._nbins)), 0)
          x = tf.reshape(x, out_shape)
          return tf.math.softmax(x, axis=-1) * (2 - self._nbins * 1e-2) + 1e-2

        def _slopes(x):
          out_shape = tf.concat((
            tf.shape(x)[:-1], (nunits, self._nbins - 1)), 0)
          x = tf.reshape(x, out_shape)
          return tf.math.softplus(x) + 1e-2

        self._bin_widths = tf.keras.layers.Dense(
            nunits * self._nbins, activation=_bin_positions, name='w')
        self._bin_heights = tf.keras.layers.Dense(
            nunits * self._nbins, activation=_bin_positions, name='h')
        self._knot_slopes = tf.keras.layers.Dense(
            nunits * (self._nbins - 1), activation=_slopes, name='s')
        self._built = True

      return tfb.RationalQuadraticSpline(
          bin_widths=self._bin_widths(x),
          bin_heights=self._bin_heights(x),
          knot_slopes=self._knot_slopes(x))

  xs = np.random.randn(3, 15).astype(np.float32)  # Keras won't Dense(.)(vec).
  splines = [SplineParams() for _ in range(nsplits)]

  def spline_flow():
    stack = tfb.Identity()
    for i in range(nsplits):
      stack = tfb.RealNVP(5 * i, bijector_fn=splines[i])(stack)
    return stack

  ys = spline_flow().forward(xs)
  ys_inv = spline_flow().inverse(ys)  # ys_inv ~= xs
  ```

  For a one-at-a-time autoregressive flow as in [1], it would be profitable to
  implement a mask over `xs` to parallelize either the inverse or the forward
  pass and implement the other using a `tf.while_loop`. See
  `tfp.bijectors.MaskedAutoregressiveFlow` for support doing so (paired with
  `tfp.bijectors.Invert` depending which direction should be parallel).

  #### References

  [1]: <NAME>, <NAME>, <NAME>, <NAME>. Neural
       Spline Flows. _arXiv preprint arXiv:1906.04032_, 2019.
       https://arxiv.org/abs/1906.04032
  """<commit_msg>modified SplineParams example for improved clarity<commit_after># Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Piecewise Rational Quadratic Spline bijector."""

import collections
import functools

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.bijectors import bijector
from tensorflow_probability.python.bijectors import softplus as softplus_bijector
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import parameter_properties
from tensorflow_probability.python.internal import tensor_util
from tensorflow_probability.python.internal import tensorshape_util


def _ensure_at_least_1d(t):
  t = tf.convert_to_tensor(t)
  return t + tf.zeros([1], dtype=t.dtype)


def _padded(t, lhs, rhs=None):
  """Left pads and optionally right pads the innermost axis of `t`."""
  lhs = tf.convert_to_tensor(lhs, dtype=t.dtype)
  zeros = tf.zeros([tf.rank(t) - 1, 2], dtype=tf.int32)
  lhs_paddings = tf.concat([zeros, [[1, 0]]], axis=0)
  result = tf.pad(t, paddings=lhs_paddings, constant_values=lhs)
  if rhs is not None:
    rhs = tf.convert_to_tensor(rhs, dtype=t.dtype)
    rhs_paddings = tf.concat([zeros, [[0, 1]]], axis=0)
    result = tf.pad(result, paddings=rhs_paddings, constant_values=rhs)
  return result


def _knot_positions(bin_sizes, range_min):
  return _padded(tf.cumsum(bin_sizes, axis=-1) + range_min, lhs=range_min)


_SplineShared = collections.namedtuple(
    'SplineShared', 'out_of_bounds,x_k,y_k,d_k,d_kp1,h_k,w_k,s_k')


class RationalQuadraticSpline(bijector.AutoCompositeTensorBijector):
  """A piecewise rational quadratic spline, as developed in [1].

  This transformation represents a monotonically increasing piecewise rational
  quadratic function. Outside of the bounds of `knot_x`/`knot_y`, the transform
  behaves as an identity function.

  Typically this bijector will be used as part of a chain, with splines for
  trailing `x` dimensions conditioned on some of the earlier `x` dimensions, and
  with the inverse then solved first for unconditioned dimensions, then using
  conditioning derived from those inverses, and so forth. For example, if we
  split a 15-D `xs` vector into 3 components, we may implement a forward and
  inverse as follows:

  ```python
  nsplits = 3

  class SplineParams(tf.Module):

    def __init__(self, nbins=32, interval_width=2, range_min=-1,
                 min_bin_width=1e-3, min_slope=1e-3):
      self._nbins = nbins
      self._interval_width = interval_width
      self._range_min = range_min
      self._min_bin_width = min_bin_width
      self._min_slope = min_slope
      self._built = False
      self._bin_widths = None
      self._bin_heights = None
      self._knot_slopes = None

    def __call__(self, x, nunits):
      if not self._built:
        def _bin_positions(x):
          out_shape = tf.concat((tf.shape(x)[:-1], (nunits, self._nbins)), 0)
          x = tf.reshape(x, out_shape)
          return tf.math.softmax(x, axis=-1) * (
                self._interval_width - self._nbins * self._min_bin_width) + \
                 self._min_bin_width

        def _slopes(x):
          out_shape = tf.concat((
            tf.shape(x)[:-1], (nunits, self._nbins - 1)), 0)
          x = tf.reshape(x, out_shape)
          return tf.math.softplus(x) + self._min_slope

        self._bin_widths = tf.keras.layers.Dense(
          nunits * self._nbins, activation=_bin_positions, name='w')
        self._bin_heights = tf.keras.layers.Dense(
          nunits * self._nbins, activation=_bin_positions, name='h')
        self._knot_slopes = tf.keras.layers.Dense(
          nunits * (self._nbins - 1), activation=_slopes, name='s')
        self._built = True

      return tfb.RationalQuadraticSpline(
        bin_widths=self._bin_widths(x),
        bin_heights=self._bin_heights(x),
        knot_slopes=self._knot_slopes(x),
        range_min=self._range_min)

  xs = np.random.randn(3, 15).astype(np.float32)  # Keras won't Dense(.)(vec).
  splines = [SplineParams() for _ in range(nsplits)]

  def spline_flow():
    stack = tfb.Identity()
    for i in range(nsplits):
      stack = tfb.RealNVP(5 * i, bijector_fn=splines[i])(stack)
    return stack

  ys = spline_flow().forward(xs)
  ys_inv = spline_flow().inverse(ys)  # ys_inv ~= xs
  ```

  For a one-at-a-time autoregressive flow as in [1], it would be profitable to
  implement a mask over `xs` to parallelize either the inverse or the forward
  pass and implement the other using a `tf.while_loop`. See
  `tfp.bijectors.MaskedAutoregressiveFlow` for support doing so (paired with
  `tfp.bijectors.Invert` depending which direction should be parallel).

  #### References

  [1]: <NAME>, <NAME>, <NAME>, <NAME>. Neural
       Spline Flows. _arXiv preprint arXiv:1906.04032_, 2019.
       https://arxiv.org/abs/1906.04032
  """
{'fraction_non_alphanumeric': 0.1155538908137466, 'fraction_numerical': 0.013765333832755876, 'mean_word_length': 3.774251229324989, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1793429', 'n_tokens_mistral': 3641, 'n_tokens_neox': 3443, 'n_words': 970}