'keras.Model ValueError/TypeError when output is from customized layer

I'm attempting to model a function as a Keras Functional model combined with a customized Keras Layer class. The idea is to simply have the Keras layer's call method use a predefined function. The function will take as input a given tensor, say tensor1, and compute a certain dot product for every value t in tensor1. The dot product relies on tensors (vectors) of different lengths, but the output is a tensor (vector) of the same length as tensor1.

When passing tensor1 into the function TheFunction(tensor) (defined below), it returns the output I expected. Moreover, when building the TheFunctionLayer object and passing in tensor1, the output is also as expected, being the same as that of TheFunction(tensor1).

The issue arises when I try to use Keras's Functional model on the same thing. I'm trying to have it so that I can build the model object and pass tensor1 into it. However, when attempting to do build the object, I get a ValueError regarding the dimensions of the tensors at play in the dot product. But if I make one small change to the output declaration in the definition of the model, I get a TypeError. Below is my source code and the issues that happen as they arise:

Imports:

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.layers import Input, Layer
import numpy as np

Tensors involved:

tensor1 = [1,2,3,4,5,6,7,8,9,10]
tensor1 = tf.convert_to_tensor(tensor1, dtype=tf.float32)

tensor2 = [1,2,3,4,5]
tensor2 = tf.convert_to_tensor(tensor2, dtype=tf.float32)

tensor3 = [6,7,8,9,10]
tensor3 = tf.convert_to_tensor(tensor3, dtype=tf.float32)

The dot product function:

def TheFunction(tensor):
    tensor = tf.convert_to_tensor(tensor)
    def prod_term(t):
        return tensor3*t
    
    def dot(t):
        return tf.tensordot(tensor2,prod_term(t),axes=1)
    
    return tf.map_fn(lambda t: dot(t), tensor)

Result of into the dot product function (no issues with dimensions):

<tf.Tensor: shape=(10,), dtype=float32, numpy=
array([ 130.,  260.,  390.,  520.,  650.,  780.,  910., 1040., 1170.,
       1300.], dtype=float32)>

The customized Keras layer (only call, no __init__ needed):

class TheFunctionLayer(Layer):
    
    def call(self,tensor):
        return TheFunction(tensor)

Result of TheFunctionLayer()(tensor1) (again, no issues with dimensions):

<tf.Tensor: shape=(10,), dtype=float32, numpy=
array([ 130.,  260.,  390.,  520.,  650.,  780.,  910., 1040., 1170.,
       1300.], dtype=float32)>

The Keras Functional model:

def TheFunctionModel(input_shape):
    x = Input(shape=input_shape)
    y = TheFunctionLayer()(x)
    
    model = keras.Model(inputs=x, outputs=y)
    
    return model

Result of building model object TheFunctionModel(tf.shape(tensor1)) (conflict with dimensions):

ValueError: Dimensions must be equal, but are 5 and 10 for '{{node the_function_layer_19/map/while/mul}} = Mul[T=DT_FLOAT](the_function_layer_19/map/while/mul/x, the_function_layer_19/map/while/TensorArrayV2Read/TensorListGetItem)' with input shapes: [5], [10].

What can possibly be going on here? I'm still new to the way TensorFlow and Keras works, but it seems like if the the output of my function works and has no conflict with tensor dimensions, why does trying to build the model object give me this dimensionality issue?

The other thing I'll add is the TypeError I get when changing the output y = TheFunctionLayer()(x) in the definition of TheFunctionModel to y = TheFunctionLayer().call(x):

def TheFunctionModel(input_shape):
    x = Input(shape=input_shape)
    y = TheFunctionLayer().call(x)
    
    model = keras.Model(inputs=x, outputs=y)
    
    return model

Result of TheFunctionModel(tf.shape(tensor1)) (TypeError occurs instead of original ValueError):

TypeError: Could not build a TypeSpec for <KerasTensor: shape=(None, 10) dtype=float32 (created by layer 'tf.convert_to_tensor')> with type KerasTensor

If changing the output in the model definition is the way to get rid of the value error, then I'm not sure I understand the type error I get and how to fix it.

Suggestions and or solutions to the first value error or even the second type error are very much appreciated.



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source