Using `exec()` in Layer in Model, `exec()` doesn't recognize any values | TensorFlow

Question

I am trying to define and use a custom layer I made using tf.keras. For it, I need to use the exec() function (you can see more in this post). The layer works, however, when trying to use it in a Keras model with the Sequential or Functional API (or when trying to call .fit() with the Subclassing API), I get an error about the exec() function, basically saying that everything I try to use in it is undefined.

import tensorflow as tf
from tensorflow import keras

def neuron(x, W, b):
    return tf.add(tf.matmul(x, W), b)

class Layer(keras.layers.Layer):
    def __init__(self, n_units=5):
        super(Layer, self).__init__()

        self.n_units = n_units

    def build(self, input_shape):
        for i in range(self.n_units):
            exec(f'self.kernel_{i} = self.add_weight("kernel_{i}", shape=[int(input_shape[1]), 1])')
            exec(f'self.bias_{i} = self.add_weight("bias_{i}", shape=[1, 1])')

    def call(self, inputs):
        for i in range(self.n_units):
            exec(f'out_{i} = neuron(inputs, self.kernel_{i}, self.bias_{i})')
        return eval(f'tf.concat([{", ".join([ f"out_{i}" for i in range(self.n_units) ])}], axis=1)')

As you can see, I'm using exec calls to dynamically create the code for however many neurons I want in the layer.

layer = Layer(4)
layer(tf.ones([3, 5]))

<tf.Tensor: shape=(3, 4), dtype=float32, numpy=
array([[ 0.13784659, -0.39448535,  1.2321193 , -1.4131087 ],
       [ 0.13784659, -0.39448535,  1.2321193 , -1.4131087 ],
       [ 0.13784659, -0.39448535,  1.2321193 , -1.4131087 ]],
      dtype=float32)>

The layer works, it outputs the correct shapes, has the trainable parameters, and TensorFlow can get it's gradients. However, when I try to incorperate this into a model this happens:

model = keras.Sequential([
    keras.layers.Dense(5, input_shape=(6,)),
    Layer(n_units=4),
    keras.layers.Dense(3)
])

NameError: in user code:

    <ipython-input-3-4c83a9a58db6>:15 call  *
        exec(f'out_{i} = neuron(inputs, self.kernel_{i}, self.bias_{i})')
    <string>:1 <module>  **
        

    NameError: name 'neuron' is not defined

Here's a pastebin of the full traceback.

The things I tried:

Define the neuron() function inside of the Layer class, and call it up using self.neuron. Gave the error NameError: name 'self' is not defined.
Define neuron within the call() function itself. Same error
Use the Function or the Subclassing API instead. Functional API was the same, and the Subclassing API gave the error when trying to call .fit() on the model instead.

I have also found this other post, showing the usage of exec() within a function. While I think the answer to that post is a step in the right direction, it doesn't fully answer my question here. In that post, the exec() function wouldn't return any value, while here it doesn't recognize any value.
Finally, I'm aware of the probable solution of using the Subclassing API and using my own custom training loop. However, I would rather solve this problem than run away from it :)

Using the exec() function in a Keras layer, when calling that Layer in a Keras Model, exec() can't seem to recognize any other values; How do I fix this?
Thanks

EDIT 1
For the traceback of the error I'm getting (pastebin here), I've located it in the tensorflow source code:
tensorflow/python/training/tracking/base.py Line 530
tensorflow/python/keras/engine/sequential.py Line 141
tensorflow/python/training/tracking/base.py Line 530
tensorflow/python/keras/engine/sequential.py Line 228
tensorflow/python/keras/engine/base_layer.py Line 997
tensorflow/python/keras/engine/base_layer.py Line 1135
tensorflow/python/keras/engine/base_layer.py Line 867
tensorflow/python/keras/engine/base_layer.py Line 907
tensorflow/python/autograph/impl/api.py Line 695

Using `exec()` in Layer in Model, `exec()` doesn't recognize any values | TensorFlow

0 Answers0