"TypeError: Could not build a TypeSpec for name ..." when using tf.assert on unknown dimen-CodePudding

Suppose I am building a model with partially known dimensions:

# Multilayer Perceptron
from keras.layers import Input
from keras.layers import Dense
import tensorflow as tf

inputs = Input(shape=(10,))

hidden1 = Dense(10, activation='relu')(inputs)
hidden2 = tf.math.add(hidden1, 5)
hidden3 = tf.math.add(hidden2, 5)
hidden4 = my_custom_op(hidden2, hidden3)
output = Dense(10, activation='linear')(hidden4)

my_custom_op is a large, complex tensor function that uses assertions in various places to ensure that assumptions about shape and rank are being met. For the sake of reproducing the problem I'll just make it this:

def my_custom_op(hidden_x, hidden_y):
    tf.assert_equal(tf.shape(hidden_x), tf.shape(hidden_y))    
    return hidden

When I run this I get the following error:

TypeError: Could not build a TypeSpec for name: "tf.debugging.assert_equal_1/assert_equal_1/Assert/Assert"

I don't understand what this error message is telling me. If I run tf.assert_equal(2, 2) I don't get an exception, so I assume it's to do with the fact that the dimensions are not known yet.

But isn't the point of asserts that they run when the dimensions are known? If not, does this mean I can't use asserts in my_custom_op because they will cause these errors when the graph is constructed?

Here is the full error message:

TypeError: Could not build a TypeSpec for name: "tf.debugging.assert_equal_1/assert_equal_1/Assert/Assert"
op: "Assert"
input: "tf.debugging.assert_equal_1/assert_equal_1/All"
input: "tf.debugging.assert_equal_1/assert_equal_1/Assert/Assert/data_0"
input: "tf.debugging.assert_equal_1/assert_equal_1/Assert/Assert/data_1"
input: "Shape"
input: "tf.debugging.assert_equal_1/assert_equal_1/Assert/Assert/data_3"
input: "Shape_1"
attr {
  key: "T"
  value {
    list {
      type: DT_STRING
      type: DT_STRING
      type: DT_INT32
      type: DT_STRING
      type: DT_INT32
    }
  }
}
attr {
  key: "summarize"
  value {
    i: 3
  }
}
 of unsupported type <class 'tensorflow.python.framework.ops.Operation'>.

CodePudding user response：

The problem is you cannot feed Keras symbolic tensors to certain Tensorflow APIs. Just wrap your function my_custom_op in a Lambda or custom layer and it should work:

import tensorflow as tf

def my_custom_op(x):
    hidden_x, hidden_y = x
    tf.assert_equal(tf.shape(hidden_x), tf.shape(hidden_y))    
    return hidden_x

inputs = tf.keras.layers.Input(shape=(10,))

hidden1 = tf.keras.layers.Dense(10, activation='relu')(inputs)
hidden2 = tf.math.add(hidden1, 5)
hidden3 = tf.math.add(hidden2, 5)
hidden4 = tf.keras.layers.Lambda(my_custom_op)([hidden2, hidden3])
output = tf.keras.layers.Dense(10, activation='linear')(hidden4)