def __new__(mcs, name, bases, attrs):
operators = set(tf.Tensor.OVERLOADABLE_OPERATORS)
operators.difference_update({'__eq__', '__ne__'})
operators.update({'__iter__'})
attrs.update((attr, _wrap_method(attr)) for attr in operators)
# Support methods for __iter__ and __bool__
private_methods = {
name
for name in dir(tf.Tensor) if name.startswith('_disallow')
}
attrs.update((attr, _wrap_method(attr)) for attr in private_methods)
if JAX_MODE or NUMPY_MODE:
other_attrs = {'__array_priority__'}
if six.PY2:
other_attrs.add('__nonzero__')
else:
other_attrs.add('__bool__')
attrs.update(
(attr, getattr(np.ndarray, attr)) for attr in other_attrs)
else:
attrs.update(
(attr, getattr(tf.Tensor, attr))
for attr in {'__bool__', '__array_priority__', '__nonzero__'})
cls = super(TensorMetaClass, mcs).__new__(mcs, name, bases, attrs)
tf.register_tensor_conversion_function(cls, conversion_func=_tensorize)
return cls
def test_works_with_registered(self):
class CustomClass:
def value(self):
return tf.convert_to_tensor(42.)
tf.register_tensor_conversion_function(
CustomClass, lambda value, **_: value.value())
tf_utils.register_symbolic_tensor_type(CustomClass)
if tf.executing_eagerly():
self.assertFalse(tf_utils.is_symbolic_tensor(
tf.Variable(name='blah', initial_value=0.)))
self.assertFalse(
tf_utils.is_symbolic_tensor(
tf.convert_to_tensor(0.)))
self.assertFalse(tf_utils.is_symbolic_tensor(
tf.SparseTensor(
indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])))
self.assertFalse(tf_utils.is_symbolic_tensor(CustomClass()))
else:
self.assertTrue(tf_utils.is_symbolic_tensor(
tf.Variable(name='blah', initial_value=0.)))
self.assertTrue(
tf_utils.is_symbolic_tensor(
tf.convert_to_tensor(0.)))
self.assertTrue(tf_utils.is_symbolic_tensor(
tf.SparseTensor(
indices=[[0, 0], [1, 2]], values=[1, 2], dense_shape=[3, 4])))
self.assertTrue(tf_utils.is_symbolic_tensor(CustomClass()))
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
"""Defer an operator overload to `tf.Tensor`.
We pull the operator out of tf.Tensor dynamically to avoid ordering issues.
Args:
cls: Class to overload operator.
op: Python string representing the operator name.
"""
@functools.wraps(getattr(tf.Tensor, op))
def _run_op(a, *args):
return getattr(tf.Tensor, op)(a.value, *args)
setattr(cls, op, _run_op)
def _tensor_conversion_function(v, dtype=None, name=None, as_ref=False):
del name, as_ref # unused
if dtype and not dtype.is_compatible_with(v.dtype):
raise ValueError(
"Incompatible type conversion requested to type '%s' for variable "
"of type '%s'" % (dtype.name, v.dtype.name))
return v.value
for operator in tf.Tensor.OVERLOADABLE_OPERATORS.difference(
{"__getitem__"}).union({"__iter__", "__bool__", "__nonzero__"}):
_overload_operator(RandomVariable, operator)
tf.register_tensor_conversion_function( # enable tf.convert_to_tensor
RandomVariable, _tensor_conversion_function)
def shape(self):
return (self.n_devices, ) + self.tensors[0]._shape_tuple() # pylint: disable=protected-access
@property
def dtype(self):
return self.tensors[0].dtype
def convert_sharded_tensor_to_eager_tensor(value, *args, **kwargs):
del args, kwargs
# TODO(nareshmodi): Consider a collective op to gather the tensors from the
# various devices for performance reasons.
return tf.stack(value.tensors)
tf.register_tensor_conversion_function(ShardedNdArray,
convert_sharded_tensor_to_eager_tensor)
class _PmapConfig(threading.local):
"""Simple config used to maintain state related to a current pmap call."""
def __init__(self):
super(_PmapConfig, self).__init__()
self._axis_name = None
self._devices = None
def axis_name(self):
return self._axis_name
def set_axis_name(self, axis_name):
self._axis_name = axis_name
self.assertIs(type(kernel), type(new_kernel))
if 'store_parameters_in_results' in new_kernel.parameters:
self.assertTrue(
new_kernel.parameters['store_parameters_in_results'])
def testNoParameters(self):
kernel = FakeInnerNoParameters()
new_kernel = util.enable_store_parameters_in_results(kernel)
self.assertIs(kernel, new_kernel)
class TensorConvertible(object):
pass
tf.register_tensor_conversion_function(
TensorConvertible, conversion_func=lambda *args: tf.constant(0))
class SimpleTensorWarningTest(tfp_test_util.TestCase):
# We must defer creating the TF objects until the body of the test.
# pylint: disable=unnecessary-lambda
@parameterized.parameters([lambda: tf.Variable(0)],
[lambda: tf.Variable(0)],
[lambda: TensorConvertible()])
def testWarn(self, tensor_callable):
tensor = tensor_callable()
warnings.simplefilter('always')
with warnings.catch_warnings(record=True) as triggered:
util.warn_if_parameters_are_not_simple_tensors({'a': tensor})
self.assertTrue(
class MyTuple(object):
"""Pretend user-side class for `ConvertToCompositeTensorTest ."""
def __init__(self, sequence):
self._sequence = tuple(sequence)
def __getitem__(self, key):
return self._sequence[key]
def __len__(self):
return len(self._sequence)
def __iter__(self):
return iter(self._sequence)
tf.register_tensor_conversion_function(
MyTuple, conversion_func=lambda x, *_, **__: tensor_tuple.TensorTuple(x))
@test_util.test_all_tf_execution_regimes
class CustomConvertToCompositeTensorTest(test_util.TestCase):
def test_iter(self):
x = MyTuple((1, [2., 3.], [[4, 5], [6, 7]]))
y = ops.convert_to_tensor_or_composite(value=x)
# TODO(jsimsa): The behavior of `is_tensor` for composite tensors have
# changed (from True to False) and this check needs to be disabled so that
# both TAP presubmits (running at HEAD) and Kokoro presubmit (using TF
# nightly) pass. Re-enable this check when TF nightly picks up this change.
# self.assertTrue(tf.is_tensor(y))
self.assertIsInstance(y, tensor_tuple.TensorTuple)
self.assertLen(y, 3)
for x_, y_ in zip(x, y):
def test_enables_nontensor_plumbing(self):
if tf.executing_eagerly():
self.skipTest('`compile` functionality changed.')
# Setup.
class Foo:
def __init__(self, input_):
self._input = input_
self.value = tf.convert_to_tensor([[42.]])
@property
def dtype(self):
return self.value.dtype
tf.register_tensor_conversion_function(
Foo, lambda x, *args, **kwargs: x.value)
tf_utils.register_symbolic_tensor_type(Foo)
class PlumbingLayer(keras.layers.Lambda):
def __init__(self, fn, **kwargs):
def _fn(*fargs, **fkwargs):
d = fn(*fargs, **fkwargs)
x = tf.convert_to_tensor(d)
d.shape = x.shape
d.get_shape = x.get_shape
return d, x
super(PlumbingLayer, self).__init__(_fn, **kwargs)
self._enter_dunder_call = False
def __call__(self, inputs, *args, **kwargs):
self._enter_dunder_call = True
d, _ = super(PlumbingLayer, self).__call__(inputs, *args, **kwargs)
self._enter_dunder_call = False
return d
def call(self, inputs, *args, **kwargs):
d, v = super(PlumbingLayer, self).call(inputs, *args, **kwargs)
if self._enter_dunder_call:
return d, v
return d
# User-land.
model = keras.Sequential([
keras.layers.InputLayer((1,)),
PlumbingLayer(Foo), # Makes a `Foo` object.
])
# Let's ensure Keras graph history is preserved by composing the models.
model = keras.Model(model.inputs, model(model.outputs))
# Now we instantiate the model and verify we have a `Foo` object, not a
# `Tensor`.
y = model(tf.convert_to_tensor([[7.]]))
self.assertIsInstance(y, Foo)
# Confirm that (custom) loss sees `Foo` instance, not Tensor.
obtained_prediction_box = [None]
def custom_loss(y_obs, y_pred):
del y_obs
obtained_prediction_box[0] = y_pred
return y_pred
# Apparently `compile` calls the loss function enough to trigger the
# side-effect.
model.compile('SGD', loss=custom_loss)
self.assertIsInstance(obtained_prediction_box[0], Foo)
def __init__(self, x, name=None):
super(FakeModule, self).__init__(name)
self.x = tensor_util.convert_nonref_to_tensor(x)
@property
def dtype(self):
return tf.as_dtype(self.x.dtype)
@property
def shape(self):
return tf.TensorShape(self.x.shape)
tf.register_tensor_conversion_function(
base_type=FakeModule,
conversion_func=lambda d, *_, **__: tf.convert_to_tensor(d.x))
@test_util.test_all_tf_execution_regimes
class ConvertNonrefToTensorTest(test_util.TestCase):
def test_np_object(self):
x = np.array(0.)
y = tensor_util.convert_nonref_to_tensor(x)
self.assertIsInstance(y, tf.Tensor)
self.assertEqual(x, self.evaluate(y))
def test_tf_tensor(self):
x = tf.constant(1.)
y = tensor_util.convert_nonref_to_tensor(x)
return self.read_value().__matmul__(o)
except AttributeError:
# See https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
def __rmatmul__(self, o):
try:
return self.read_value().__rmatmul__(o)
except AttributeError:
# See https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
# pylint: enable=multiple-statements
tf.register_tensor_conversion_function(AutoCastVariable,
AutoCastVariable._dense_var_to_tensor) # pylint:disable=protected-access
def create_autocast_variable(variable):
"""Creates an AutoCastVariable that wraps another variable.
This typically just returns `AutoCastVariable(variable)`. But, if the variable
is a DistributedVariable or one of its subclasses, we instead dynamically
create a class that subclasses from both AutoCastVariable and
variable.__class__. This is so the returned variable will still pass
`isinstance(variable, variable.__class__)`, which is required for
DistributedVariables and its subclasses to work properly.
Args:
variable: A floating-point resource variable to wrap.
return self.read_value().__matmul__(o)
except AttributeError:
# See
# https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
def __rmatmul__(self, o):
try:
return self.read_value().__rmatmul__(o)
except AttributeError:
# See
# https://docs.python.org/3/library/constants.html#NotImplemented
return NotImplemented
tf.register_tensor_conversion_function(AutoCastVariable,
AutoCastVariable._dense_var_to_tensor)
def create_autocast_variable(variable):
"""Creates an AutoCastVariable that wraps another variable.
This typically just returns `AutoCastVariable(variable)`. But, if the
variable is a DistributedVariable or one of its subclasses, we instead
dynamically create a class that subclasses from both AutoCastVariable and
variable.__class__. This is so the returned variable will still pass
`isinstance(variable, variable.__class__)`, which is required for
DistributedVariables and its subclasses to work properly.
Args:
variable: A floating-point resource variable to wrap.
@property
def shape(self):
return self.data.shape
@property
def dtype(self):
return self.data.dtype
def fail_on_convert(x, **kwargs):
_ = x
_ = kwargs
raise TypeError('Cannot convert DummyArrayLike to a tensor')
tf.register_tensor_conversion_function(DummyArrayLike, fail_on_convert)
class DataAdapterTestBase(keras_parameterized.TestCase):
def setUp(self):
super(DataAdapterTestBase, self).setUp()
self.batch_size = 5
self.numpy_input = np.zeros((50, 10))
self.numpy_target = np.ones(50)
self.tensor_input = tf.constant(2.0, shape=(50, 10))
self.tensor_target = tf.ones((50, ))
self.arraylike_input = DummyArrayLike(self.numpy_input)
self.arraylike_target = DummyArrayLike(self.numpy_target)
self.dataset_input = tf.data.Dataset.from_tensor_slices(
(self.numpy_input,
self.numpy_target)).shuffle(50).batch(self.batch_size)
