<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.
# ============================================================================
"""The `DeferredTensor` class."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import six

import tensorflow.compat.v2 as tf
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import name_util
from tensorflow_probability.python.internal import tensorshape_util


__all__ = [
    'DeferredTensor',
]


def _wrap_method(cls, attr):
  """Replaces member function's first arg, `self`, to `self._value()`.

  This function is used by `_get_tensor_like_attributes` to take existing
  `Tensor` member functions and make them operate on `self._value()`, i.e., the
  concretization of a `Distribution`.

  Args:
    cls: The `class` from which we will look up the `attr`.
    attr: Python `str` representing the `attr` to inject a new notion of `self`.

  Returns:
    dependency_injected_function: Python `callable` (or `property`)
      corresponding to `cls.attr` with `self` replaced as `self._value()`.
  """
  fn = getattr(cls, attr)
  is_property = isinstance(fn, property)
  if is_property:
    fn = fn.fget
  @functools.wraps(fn)
  def wrapped(self, *args, **kwargs):
    return fn(self._value(), *args, **kwargs)  # pylint: disable=protected-access
  return property(wrapped) if is_property else wrapped


def _tensorize(d, dtype=None, name=None, as_ref=False):
  """Tensor conversion function presuming `hasattr(d, '_value')`."""
  return d._value(dtype, name, as_ref)  # pylint: disable=protected-access


class TensorMetaClass(type):
  """A type of class which will make objects which act like Tensors."""

  def __new__(mcs, name, bases, attrs):
    attrs.update(
        (attr, _wrap_method(tf.Tensor, attr))
        for attr in tf.Tensor.OVERLOADABLE_OPERATORS.union({'__iter__'}))
    attrs.update(
        (attr, getattr(tf.Tensor, attr))
        for attr in {'__nonzero__', '__bool__', '__array_priority__'})
    cls = super(TensorMetaClass, mcs).__new__(mcs, name, bases, attrs)
    tf.register_tensor_conversion_function(cls, conversion_func=_tensorize)
    return cls


NONE_SPECIFIED = 'None'


@six.add_metaclass(TensorMetaClass)
class DeferredTensor(tf.Module):
  """Variable tracking object which applies function upon `convert_to_tensor`.

  #### Example

  ```python
  import tensorflow.compat.v2 as tf

  trainable_normal = tfd.Normal(
      loc=tf.Variable(0.),
      scale=tfp.util.DeferredTensor(tf.math.exp, tf.Variable(0.)))

  trainable_normal.loc
  # ==> <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.0>

  trainable_normal.scale
  # ==> <DeferredTensor: dtype=float32, shape=[], fn=exp>

  # Operators work with `DeferredTensor`.<commit_msg>Fix failure due to tensor equality changes<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.
# ============================================================================
"""The `DeferredTensor` class."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import six

import tensorflow.compat.v2 as tf
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import name_util
from tensorflow_probability.python.internal import tensorshape_util


__all__ = [
    'DeferredTensor',
]


def _wrap_method(cls, attr):
  """Replaces member function's first arg, `self`, to `self._value()`.

  This function is used by `_get_tensor_like_attributes` to take existing
  `Tensor` member functions and make them operate on `self._value()`, i.e., the
  concretization of a `Distribution`.

  Args:
    cls: The `class` from which we will look up the `attr`.
    attr: Python `str` representing the `attr` to inject a new notion of `self`.

  Returns:
    dependency_injected_function: Python `callable` (or `property`)
      corresponding to `cls.attr` with `self` replaced as `self._value()`.
  """
  fn = getattr(cls, attr)
  is_property = isinstance(fn, property)
  if is_property:
    fn = fn.fget
  @functools.wraps(fn)
  def wrapped(self, *args, **kwargs):
    return fn(self._value(), *args, **kwargs)  # pylint: disable=protected-access
  return property(wrapped) if is_property else wrapped


def _tensorize(d, dtype=None, name=None, as_ref=False):
  """Tensor conversion function presuming `hasattr(d, '_value')`."""
  return d._value(dtype, name, as_ref)  # pylint: disable=protected-access


class TensorMetaClass(type):
  """A type of class which will make objects which act like Tensors."""

  def __new__(mcs, name, bases, attrs):
    operators = tf.Tensor.OVERLOADABLE_OPERATORS
    operators.difference_update({'__eq__', '__ne__'})
    operators.update({'__iter__'})
    attrs.update((attr, _wrap_method(tf.Tensor, attr)) for attr in operators)
    attrs.update(
        (attr, getattr(tf.Tensor, attr))
        for attr in {'__nonzero__', '__bool__', '__array_priority__'})
    cls = super(TensorMetaClass, mcs).__new__(mcs, name, bases, attrs)
    tf.register_tensor_conversion_function(cls, conversion_func=_tensorize)
    return cls


NONE_SPECIFIED = 'None'


@six.add_metaclass(TensorMetaClass)
class DeferredTensor(tf.Module):
  """Variable tracking object which applies function upon `convert_to_tensor`.

  #### Example

  ```python
  import tensorflow.compat.v2 as tf

  trainable_normal = tfd.Normal(
      loc=tf.Variable(0.),
      scale=tfp.util.DeferredTensor(tf.math.exp, tf.Variable(0.)))

  trainable_normal.loc
  # ==> <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=0.0>

  trainable_normal.scale
  # ==> <DeferredTensor: dtype=float32, shape=[], fn=exp>

  # Operators work with `DeferredTensor`.
{'fraction_non_alphanumeric': 0.12801270024534564, 'fraction_numerical': 0.005484196853802858, 'mean_word_length': 4.694330320460148, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 11, 'https://': 0, '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': '822668', 'n_tokens_mistral': 2101, 'n_tokens_neox': 2042, 'n_words': 661}