Issue 1:
System information

Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colaboratory

TensorFlow installed from (source or binary):

TensorFlow version (use command below): 1.9.0-rc2

Python version: version 3.6

Bazel version (if compiling from source):

GCC/Compiler version (if compiling from source):

CUDA/cuDNN version:

GPU model and memory: Colab GPU

Exact command to reproduce: https://gist.github.com/nairouz/035a830d1e58a3759a6a1e193f5defed

Describe the problem

I have a simple Keras model with one custom layer which works fine on the graph based execution. When I switched to eager execution via tf.enable_eager_execution(), I got stuck on a weird error.

Source code

import numpy as np
import tensorflow as tf
import tensorflow.keras.backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer, Input
from tensorflow.keras.losses import kullback_leibler_divergence

tf.enable_eager_execution()

class ClusteringLayer(Layer):
def __init__(self, output_dim, input_dim=None, alpha=1.0, **kwargs):
self.output_dim = output_dim
self.input_dim = input_dim
self.alpha = alpha
super(ClusteringLayer, self).__init__(**kwargs)

def build(self, input_shape):
self.W = self.add_weight(name='kernel', shape=(self.output_dim, input_shape[1].value), initializer='Identity', trainable=True)
super(ClusteringLayer, self).build(input_shape)

def call(self, x, mask=None):
q = 1.0/(1.0 + K.sqrt(K.sum(K.square(K.expand_dims(x, 1) - self.W), axis=2))**2 /self.alpha)
q = q**((self.alpha+1.0)/2.0)
q = K.transpose(K.transpose(q)/K.sum(q, axis=1))
return q

def compute_output_shape(self, input_shape):
return (input_shape[0], self.output_dim)

def clustering_loss(y_true, y_pred):
a = K.square(y_pred) / K.sum(y_pred, axis=0)
p = K.transpose(K.transpose(a) / K.sum(a, axis=1))
loss = kullback_leibler_divergence(p, y_pred)
return loss

input1 = Input(shape=(10,), name="input")
out = ClusteringLayer(output_dim = 5, name='clustering')(input1)
model = Model(inputs=input1, outputs=out)
model.compile(optimizer=tf.train.AdamOptimizer(1e-3), loss={'clustering' : clustering_loss})
np.random.seed(0)
X = np.random.random((20, 10)).astype(np.float32)
Y = np.random.random((20, 5)).astype(np.float32)
model.fit(x={'input' : X}, y={'clustering' : Y}, batch_size=1, epochs=10)

###Logs

clustering

Epoch 1/10

AssertionError Traceback (most recent call last)

in ()

37 Y = np.random.random((20, 5))

38

---> 39 model.fit(x={'input' : X}, y={'clustering' : Y}, batch_size=1, epochs=10)

40

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, **kwargs)

1331 initial_epoch=initial_epoch,

1332 steps_per_epoch=steps_per_epoch,

-> 1333 validation_steps=validation_steps)

1334 else:

1335 return training_arrays.fit_loop(

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_eager.py in fit_loop(model, inputs, targets, sample_weights, class_weight, val_inputs, val_targets, val_sample_weights, batch_size, epochs, verbose, callbacks, shuffle, callback_metrics, initial_epoch, steps_per_epoch, validation_steps)

1039 shuffle=shuffle,

1040 num_train_samples=num_train_samples,

-> 1041 do_validation=do_validation)

1042 callbacks.on_epoch_end(epoch, epoch_logs)

1043 if callback_model.stop_training:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_eager.py in batch_fit_loop(model, inputs, targets, epoch_logs, index_array, out_labels, callback_model, batch_size, sample_weights, val_inputs, val_targets, val_sample_weights, callbacks, shuffle, num_train_samples, do_validation)

395 targets_batch,

396 sample_weights=sample_weights_batch,

--> 397 training=True)

398

399 if not isinstance(outs, list):

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_eager.py in _process_single_batch(model, inputs, targets, sample_weights, training)

787 outs, loss, loss_metrics = _model_loss(model, inputs, targets,

788 sample_weights=sample_weights,

--> 789 training=training)

790 if loss is None:

791 raise ValueError('The model cannot be run '

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_eager.py in _model_loss(model, inputs, targets, sample_weights, training)

126 outs = model.call(inputs[0], training=training)

127 else:

--> 128 outs = model.call(inputs[0])

129 else:

130 if model._expects_training_arg:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py in call(self, inputs, training, mask)

718 outputs, _ = self._run_internal_graph(inputs,

719 training=training,

--> 720 mask=masks)

721 return outputs

722

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/network.py in _run_internal_graph(self, inputs, training, mask)

939 output_shapes = []

940 for x in self.outputs:

--> 941 assert str(id(x)) in tensor_map, 'Could not compute output ' + str(x)

942 tensor, mask = tensor_map[str(id(x))]

943 output_shapes.append(backend.int_shape(x))

AssertionError: Could not compute output DeferredTensor('None', shape=(5,), dtype=float32)

Issue 2:
When an op constructing function is called for multiple times on the same set of inputs, each of these calls adds a new op to the graph, which leads to extra computation. I understand that this could be avoided in most cases by reusing the output tensor. However, this can't be easily done in some situations, e.g. when using different optimizers for different parts of a neural network. Therefore, IMHO, it would be beneficial to have the framework automatically merge identical ops.