Chat gpt solved it:

import tensorflow as tf
from tensorflow import keras

# Callback personalizado
class ExponentialLearningRate(keras.callbacks.Callback):
    def __init__(self, factor):
        super().__init__()
        self.factor = factor
        self.rates = []
        self.losses = []

    def on_batch_end(self, batch, logs=None):
        logs = logs or {}
        self.losses.append(logs.get("loss", None))

        # Obtener el valor de learning_rate y asegurarnos de que sea un número
        lr = self.model.optimizer.learning_rate
        if isinstance(lr, tf.Variable):
            lr = lr.numpy()  # Convertir a valor numérico si es un tf.Variable

        # Registrar la tasa de aprendizaje
        self.rates.append(lr)

        # Actualizar el learning rate
        if isinstance(lr, (float, int)):
            new_lr = lr * self.factor
            self.model.optimizer.learning_rate.assign(new_lr)  # Modificar learning_rate

# Modelo simple
model = keras.models.Sequential([
    keras.layers.Flatten(input_shape=[28, 28]),
    keras.layers.Dense(300, activation="relu"),
    keras.layers.Dense(100, activation="relu"),
    keras.layers.Dense(10, activation="softmax")
])

# Compilar modelo
model.compile(
    loss="sparse_categorical_crossentropy",
    optimizer=keras.optimizers.SGD(learning_rate=1e-3),
    metrics=["accuracy"]
)

# Crear el callback
expon_lr = ExponentialLearningRate(factor=1.005)

# Ajustar modelo
history = model.fit(
    X_train, y_train,
    epochs=1,
    validation_data=(X_valid, y_valid),
    callbacks=[expon_lr]
)