Repository: tensorflow
Issue: Unable to train/take gradient of integer variable under any condition
Labels: stat:awaiting tensorflower,type:bug,comp:ops,TF 2.16

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

v2.16.1-19-g810f233968c

### Custom code

Yes

### OS platform and distribution

Debian 11

### Mobile device

_No response_

### Python version

3.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

12.6

### GPU model and memory

_No response_

### Current behavior?

Models with integer variables cannot have gradients computed (or be trained because of this). They fail with `ValueError: No gradients provided for any variable.` Please note that I'm referring to building/training models, _not_ post-training model quantization.

For context, I'm working on a somewhat novel approach to solving a type of engineering problem that includes both discrete and continuous values. In some cases, variables must be one of a set of numeric values (e.g. `1`, `2`, `3000`, etc). This isn't something that can really be split out into separate models and trained independently, as both these variables are used throughout a complex, multiple-input multiple-output layer based model. An approximation cannot be used either, as even a highly accurate approximation can feasibly result in an incorrect result under certain conditions.

Intuitively, a derivative cannot be calculated on integer values because they're discrete, and differentiation requires a continuous function over the domain of differentiation. However, floats (and any limited-precision data type) also suffer from this - but differentiation is generally considered "viable" for floats. Under any argument made for floats, integers should also be considered differentiable. There are certainly some issues with this for small number, but they have much less of an impact for larger ones (just like floats very close to 0 vs "single digit" floats i.e. `1`, `2`, etc).

If maintainers/triagers/project folks don't want to go in this direction, users like myself should be able to implement this via a `tf.custom_gradient` function. However, even when implementing this function, gradients/training fail in the exact same manner in the exact same place. When using a custom gradient function with an integer variable, TF fails with the aforementioned error message without even calling the custom gradient function. IMO this is pretty clearly a bug. See below for specific code to reproduce this.

### Standalone code to reproduce the issue

This fails (no custom gradient function, int32 type):

```python
import keras
import tensorflow as tf
import numpy as np

variable_dtype = tf.int32
# variable_dtype = tf.float32   # Uncommenting this fixes the issue


class BugTestLayer(keras.layers.Layer):
    # Layer is just y = self.var * x
    def build(self, input_shape):
        self.var = self.add_variable(
            (1,), initializer="zeros", dtype=variable_dtype)

    def call(self, input):
        return input * self.var


input_layer = keras.Input((1,), dtype=tf.int32)
test_layer = BugTestLayer()

output = test_layer(input_layer)
model = keras.Model(inputs=[input_layer], outputs=[output])

values = np.array([[i] for i in range(1000)])

model.compile(
    loss=[keras.losses.MeanSquaredError()],
    optimizer=keras.optimizers.RMSprop(),
    metrics=[keras.metrics.MeanSquaredError()],
)

# This will always raise a `ValueError: No gradients provided for any variable.`
# when using an integer type
history = model.fit(values, values, batch_size=1, epochs=2)
```

This works (no custom gradient function, float32 type):
```python
import keras
import tensorflow as tf
import numpy as np

# variable_dtype = tf.int32
variable_dtype = tf.float32   # Uncommenting this fixes the issue


class BugTestLayer(keras.layers.Layer):
    # Layer is just y = self.var * x
    def build(self, input_shape):
        self.var = self.add_variable(
            (1,), initializer="zeros", dtype=variable_dtype)

    def call(self, input):
        return input * self.var


input_layer = keras.Input((1,), dtype=tf.int32)
test_layer = BugTestLayer()

output = test_layer(input_layer)
model = keras.Model(inputs=[input_layer], outputs=[output])

values = np.array([[i] for i in range(1000)])

model.compile(
    loss=[keras.losses.MeanSquaredError()],
    optimizer=keras.optimizers.RMSprop(),
    metrics=[keras.metrics.MeanSquaredError()],
)

# This will always raise a `ValueError: No gradients provided for any variable.`
# when using an integer type
history = model.fit(values, values, batch_size=1, epochs=2)
```

This fails (custom gradient function, int32 type):
```python
# %%

import keras
import tensorflow as tf
import numpy as np

# %%


class TestLayer(keras.layers.Layer):
    def build(self, input_shape):
        # dtype is the problem.
        # integer type results in "No gradients provided for any variable", while
        # floats work just fine.
        self.var = self.add_variable(
            (1,), initializer="zeros", dtype=tf.int32)

    @tf.custom_gradient
    def op(input, var):
        def grad(upstream, variables=None):
            return tf.squeeze(upstream * tf.cast(var, tf.float32), axis=[1]), \
                tf.squeeze(upstream * tf.cast(input, tf.float32), axis=[1])

        return tf.cast(input, tf.float32) * tf.cast(var, tf.float32), grad

    def call(self, input):
        return TestLayer.op(input, self.var)

# %%


input_layer = keras.Input((1,), name="input_layer", dtype=tf.int32)
test_layer = TestLayer()

output = test_layer(input_layer)
model = keras.Model(inputs=[input_layer], outputs=[output], name="model")

# %%

model.summary()
# Example evaluation (untrained)
model(tf.constant([[100]]))

keras.utils.plot_model(model, "my_first_model.png", show_dtype=True,
                       show_layer_names=True, show_shapes=True, show_trainable=True)


# %%

values = np.array([
    [i]
    for i in range(1000)
])

model.compile(
    loss=[keras.losses.MeanSquaredError()],
    optimizer=keras.optimizers.RMSprop(),
    metrics=[keras.metrics.MeanSquaredError()],
)

history = model.fit(values, values, batch_size=1, epochs=10)
```

This works (custom gradient function, float32 type):
```python
# %%

import keras
import tensorflow as tf
import numpy as np

# %%


class TestLayer(keras.layers.Layer):
    def build(self, input_shape):
        # dtype is the problem.
        # integer type results in "No gradients provided for any variable", while
        # floats work just fine.
        self.var = self.add_variable(
            (1,), initializer="zeros", dtype=tf.int32)

    @tf.custom_gradient
    def op(input, var):
        def grad(upstream, variables=None):
            return tf.squeeze(upstream * tf.cast(var, tf.float32), axis=[1]), \
                tf.squeeze(upstream * tf.cast(input, tf.float32), axis=[1])

        return tf.cast(input, tf.float32) * tf.cast(var, tf.float32), grad

    def call(self, input):
        return TestLayer.op(input, self.var)

# %%


input_layer = keras.Input((1,), name="input_layer", dtype=tf.float32)
test_layer = TestLayer()

output = test_layer(input_layer)
model = keras.Model(inputs=[input_layer], outputs=[output], name="model")

# %%

model.summary()
# Example evaluation (untrained)
model(tf.constant([[100]]))

keras.utils.plot_model(model, "my_first_model.png", show_dtype=True,
                       show_layer_names=True, show_shapes=True, show_trainable=True)


# %%

values = np.array([
    [i]
    for i in range(1000)
])

model.compile(
    loss=[keras.losses.MeanSquaredError()],
    optimizer=keras.optimizers.RMSprop(),
    metrics=[keras.metrics.MeanSquaredError()],
)

history = model.fit(values, values, batch_size=1, epochs=10)
```


### Relevant log output

Logs for the last failure:

```shell
Epoch 1/10
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
File /workspaces/power sim 2/test_docs.py:14
      3 values = np.array([
      4     [i]
      5     for i in range(1000)
      6 ])
      8 model.compile(
      9     loss=[keras.losses.MeanSquaredError()],
     10     optimizer=keras.optimizers.RMSprop(),
     11     metrics=[keras.metrics.MeanSquaredError()],
     12 )
---> 14 history = model.fit(values, values, batch_size=1, epochs=10)

File ~/.local/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py:122, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    119     filtered_tb = _process_traceback_frames(e.__traceback__)
    120     # To get the full stack trace, call:
    121     # `keras.config.disable_traceback_filtering()`
--> 122     raise e.with_traceback(filtered_tb) from None
    123 finally:
    124     del filtered_tb

File ~/.local/lib/python3.12/site-packages/keras/src/optimizers/base_optimizer.py:662, in BaseOptimizer._filter_empty_gradients(self, grads, vars)
    659         missing_grad_vars.append(v.name)
    661 if not filtered_grads:
--> 662     raise ValueError("No gradients provided for any variable.")
    663 if missing_grad_vars:
    664     warnings.warn(
    665         "Gradients do not exist for variables "
    666         f"{list(reversed(missing_grad_vars))} when minimizing the loss."
    667         " If using `model.compile()`, did you forget to provide a "
    668         "`loss` argument?"
    669     )

ValueError: No gradients provided for any variable.
```