<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.
# ============================================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools

# Dependency imports
import numpy as np

import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import

tfk = tf.keras
tfkl = tf.keras.layers
tfb = tfp.bijectors
tfd = tfp.distributions
tfpl = tfp.layers


def _logit_avg_expit(t):
  """Computes `logit(mean(expit(t)))` in a numerically stable manner."""
  log_avg_prob = (
      tf.reduce_logsumexp(input_tensor=-tf.nn.softplus(-t), axis=0) -
      tf.math.log(tf.cast(tf.shape(input=t)[0], t.dtype)))
  return log_avg_prob - tf.math.log1p(-tf.exp(log_avg_prob))


def _vec_pad(x, value=0):
  """Prepends a column of zeros to a matrix."""
  paddings = tf.concat(
      [tf.zeros([tf.rank(x) - 1, 2], dtype=tf.int32), [[1, 0]]], axis=0)
  return tf.pad(tensor=x, paddings=paddings, constant_values=value)


@test_util.run_all_in_graph_and_eager_modes
class EndToEndTest(tf.test.TestCase):
  """Test tfp.layers work in all three Keras APIs.

  For end-to-end tests we fit a Variational Autoencoder (VAE) because this
  requires chaining two Keras models, an encoder and decoder. Chaining two
  models is important because making a `Distribution` as output by a Keras model
  the input of another Keras model--and concurrently fitting both--is the
  primary value-add of using the `tfp.layers.DistributionLambda`. Otherwise,
  under many circumstances you can directly return a Distribution from a Keras
  layer, as long as the Distribution base class has a tensor conversion function
  registered via `tf.register_tensor_conversion_function`.

  Fundamentally, there are three ways to be Keras models:
  1. `tf.keras.Sequential`
  2. Functional API
  3. Subclass `tf.keras.Model`.

  Its important to have end-to-end tests for all three, because #1 and #2 call
  `__call__` and `call` differently. (#3's call pattern depends on user
  implementation details, but in general ends up being either #1 or #2.)
  """

  def setUp(self):
    self.encoded_size = 2
    self.input_shape = [2, 2, 1]
    self.train_size = 10000
    self.test_size = 1000
    self.x = (np.random.rand(
        self.train_size, *self.input_shape) > 0.75).astype(np.float32)
    self.x_test = (np.random.rand(
        self.test_size, *self.input_shape) > 0.75).astype(np.float32)
    super(EndToEndTest, self).setUp()

  def test_keras_sequential_api(self):
    """Test `DistributionLambda`s are composable via Keras `Sequential` API."""

    prior_model = tfk.Sequential([
        tfpl.VariableLayer(shape=[self.encoded_size]),
        tfpl.DistributionLambda(
            lambda t: tfd.Independent(tfd.Normal(loc=t, scale=1),  # pylint: disable=g-long-lambda
                                      reinterpreted_batch_ndims=1)),
    ])

    encoder_model = tfk.Sequential([
        tfkl.InputLayer(input_shape=self.input_shape),
        tfkl.Flatten(),
        tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(self.encoded_size)),
        tfpl.MultivariateNormalTriL(
            self.encoded_size,
            activity_regularizer=tfpl.KLDivergenceRegularizer(
                lambda: prior_model(0.),
                weight=0.9)),  # "beta" as in beta-VAE.
    ])

    decoder_model = tfk.Sequential([
        tfkl.InputLayer(input_shape=[self.encoded_size]),
        tfkl.Dense(tfpl.IndependentBernoulli.params_size(self.input_shape)),
        tfpl.IndependentBernoulli(self.input_shape, tfd.Bernoulli.logits),
    ])

    vae_model = tfk.Model(
        inputs=encoder_model.inputs,
        outputs=decoder_model(encoder_model.outputs[0]))
<commit_msg>Temporary workaround for b/139437503 -- set training=False when invoking Sequential layer.<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.
# ============================================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools

# Dependency imports
import numpy as np

import tensorflow as tf
import tensorflow_probability as tfp

from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import

tfk = tf.keras
tfkl = tf.keras.layers
tfb = tfp.bijectors
tfd = tfp.distributions
tfpl = tfp.layers


def _logit_avg_expit(t):
  """Computes `logit(mean(expit(t)))` in a numerically stable manner."""
  log_avg_prob = (
      tf.reduce_logsumexp(input_tensor=-tf.nn.softplus(-t), axis=0) -
      tf.math.log(tf.cast(tf.shape(input=t)[0], t.dtype)))
  return log_avg_prob - tf.math.log1p(-tf.exp(log_avg_prob))


def _vec_pad(x, value=0):
  """Prepends a column of zeros to a matrix."""
  paddings = tf.concat(
      [tf.zeros([tf.rank(x) - 1, 2], dtype=tf.int32), [[1, 0]]], axis=0)
  return tf.pad(tensor=x, paddings=paddings, constant_values=value)


@test_util.run_all_in_graph_and_eager_modes
class EndToEndTest(tf.test.TestCase):
  """Test tfp.layers work in all three Keras APIs.

  For end-to-end tests we fit a Variational Autoencoder (VAE) because this
  requires chaining two Keras models, an encoder and decoder. Chaining two
  models is important because making a `Distribution` as output by a Keras model
  the input of another Keras model--and concurrently fitting both--is the
  primary value-add of using the `tfp.layers.DistributionLambda`. Otherwise,
  under many circumstances you can directly return a Distribution from a Keras
  layer, as long as the Distribution base class has a tensor conversion function
  registered via `tf.register_tensor_conversion_function`.

  Fundamentally, there are three ways to be Keras models:
  1. `tf.keras.Sequential`
  2. Functional API
  3. Subclass `tf.keras.Model`.

  Its important to have end-to-end tests for all three, because #1 and #2 call
  `__call__` and `call` differently. (#3's call pattern depends on user
  implementation details, but in general ends up being either #1 or #2.)
  """

  def setUp(self):
    self.encoded_size = 2
    self.input_shape = [2, 2, 1]
    self.train_size = 10000
    self.test_size = 1000
    self.x = (np.random.rand(
        self.train_size, *self.input_shape) > 0.75).astype(np.float32)
    self.x_test = (np.random.rand(
        self.test_size, *self.input_shape) > 0.75).astype(np.float32)
    super(EndToEndTest, self).setUp()

  def test_keras_sequential_api(self):
    """Test `DistributionLambda`s are composable via Keras `Sequential` API."""

    prior_model = tfk.Sequential([
        tfpl.VariableLayer(shape=[self.encoded_size]),
        tfpl.DistributionLambda(
            lambda t: tfd.Independent(tfd.Normal(loc=t, scale=1),  # pylint: disable=g-long-lambda
                                      reinterpreted_batch_ndims=1)),
    ])

    encoder_model = tfk.Sequential([
        tfkl.InputLayer(input_shape=self.input_shape),
        tfkl.Flatten(),
        tfkl.Dense(tfpl.MultivariateNormalTriL.params_size(self.encoded_size)),
        tfpl.MultivariateNormalTriL(
            self.encoded_size,
            activity_regularizer=tfpl.KLDivergenceRegularizer(
                lambda: prior_model(0.),
                weight=0.9)),  # "beta" as in beta-VAE.
    ])

    decoder_model = tfk.Sequential([
        tfkl.InputLayer(input_shape=[self.encoded_size]),
        tfkl.Dense(tfpl.IndependentBernoulli.params_size(self.input_shape)),
        tfpl.IndependentBernoulli(self.input_shape, tfd.Bernoulli.logits),
    ])

    vae_model = tfk.Model(
        inputs=encoder_model.inputs,
        # TODO(b/139437503): remove training=False once cl/263432058 hits
        # nightly.
        outputs=decoder_model(encoder_model.outputs[0], training=False))

{'fraction_non_alphanumeric': 0.11566748384221083, 'fraction_numerical': 0.015489190996211277, 'mean_word_length': 3.7088142707240292, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '936922', 'n_tokens_mistral': 2933, 'n_tokens_neox': 2749, 'n_words': 785}