<commit_before>from __future__ import print_function, division, absolute_import

import tensorflow as tf
from tensorflow.python.keras.layers import Layer
from tensorflow_probability.python.distributions import Distribution
from tensorflow_probability.python.layers.internal import distribution_tensor_coercible as dtc

class ReduceMean(Layer):
  """ ReduceMean """

  def __init__(self, axis=None, keepdims=None, **kwargs):
    super(ReduceMean, self).__init__(**kwargs)
    self.axis = axis
    self.keepdims = keepdims

  def call(self, x):
    return [tf.reduce_mean(i, axis=self.axis, keepdims=self.keepdims) for i in x] \
    if isinstance(x, (tuple, list)) else \
    tf.reduce_mean(x, axis=self.axis, keepdims=self.keepdims)

class Moments(Layer):
  """ Moments """

  def __init__(self, mean=True, variance=True, **kwargs):
    super(Moments, self).__init__(**kwargs)
    self.mean = bool(mean)
    self.variance = bool(variance)<commit_msg>Sampling layer for Distribution layer<commit_after>from __future__ import print_function, division, absolute_import

import tensorflow as tf
from tensorflow.python.keras.layers import Layer
from tensorflow_probability.python.distributions import Distribution
from tensorflow_probability.python.layers.internal import distribution_tensor_coercible as dtc

class ReduceMean(Layer):
  """ ReduceMean """

  def __init__(self, axis=None, keepdims=None, **kwargs):
    super(ReduceMean, self).__init__(**kwargs)
    self.axis = axis
    self.keepdims = keepdims

  def call(self, x):
    return [tf.reduce_mean(i, axis=self.axis, keepdims=self.keepdims) for i in x] \
    if isinstance(x, (tuple, list)) else \
    tf.reduce_mean(x, axis=self.axis, keepdims=self.keepdims)

class Sampling(Layer):
  """ Sample the output from tensorflow-probability
  distribution layers """

  def __init__(self, n_samples=None, **kwargs):
    super(Sampling, self).__init__(**kwargs)
    self.n_samples = n_samples

  def call(self, x):
    assert isinstance(x, Distribution) and hasattr(x, '_keras_history'), \
    "Input to this layer must be instance of tensorflow_probability Distribution"
    if self.n_samples is not None:
      return x.sample(self.n_samples)
    else:
      return tf.convert_to_tensor(x)

class Moments(Layer):
  """ Moments """

  def __init__(self, mean=True, variance=True, **kwargs):
    super(Moments, self).__init__(**kwargs)
    self.mean = bool(mean)
    self.variance = bool(variance)
{'fraction_non_alphanumeric': 0.11274509803921569, 'fraction_numerical': 0.0, 'mean_word_length': 4.858851674641149, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, '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': '5360199', 'n_tokens_mistral': 777, 'n_tokens_neox': 771, 'n_words': 191}