Bidirectional layer rejects custom activations

[jm-willy]

Every other layer I tried accepts custom activations

Minimal code

import tensorflow as tf


@tf.function
def my_hard_sigmoid(x):
    return x


@tf.function
def my_hard_tanh(x):
    return x


inputs = tf.keras.Input(
    shape=(
        100,
        32,
    )
)
init_layer = tf.keras.layers.LSTM(
    8,
    activation=my_hard_tanh,
    recurrent_activation=my_hard_sigmoid,
)
x = tf.keras.layers.Bidirectional(init_layer)(inputs)
model = tf.keras.Model(inputs=inputs, outputs=x)
model.summary(expand_nested=True, show_trainable=True)

Output
ValueError: Unknown activation function '{'module': 'tensoions.py", line 646, in deserializerflow.python.eager.polymorphic_function.polymorphic_function', 'class_name': 'Function', 'config': 'my_hard_tanh', 'rflow.python.eager.polymorphic_function.polymorphic_function', 'class_name': 'Functregistered_name': 'Function'}' cannot be deserialized.

Versions
python 3.11.9
tensorflow 2.15.1
keras 2.15.0

[fchollet]

Your custom objects always need to be serializable. You can use @keras.saving.register_keras_serializable for this purpose https://keras.io/api/models/model_saving_apis/serialization_utils/#registerkerasserializable-function

[jm-willy]

@tf.keras.utils.register_keras_serializable("custom_activations")
@tf.function
def my_hard_sigmoid(x):
    return x


@tf.keras.utils.register_keras_serializable("custom_activations")
@tf.function
def my_hard_tanh(x):
    return x

Thanks. Doesn't works when @tf.function is added, asks for a class with a from_config() method. What a pain!
However, I using custom layers with no problems, I opened issue for anyone else with the problem.

[fchollet]

Why do you use tf.function though?

[jm-willy]

Isn't it a good speedup?