def test_works_with_registered(self):
class CustomClass(object):
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 test_enables_nontensor_plumbing(self):
if tf.executing_eagerly():
self.skipTest('`compile` functionality changed.')
# Setup.
class Foo(object):
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)