Repository: keras-team/keras (python)

Issue:
Potential error in documentation of keras.random.uniform.
In the documentation of keras.ranom.uniform I find

    seed: A Python integer or instance of
        `keras.random.SeedGenerator`.
        Used to make the behavior of the initializer
        deterministic. Note that an initializer seeded with an integer
        or None (unseeded) will produce the same random values
        across multiple calls. To get different random values
        across multiple calls, use as seed an instance
        of `keras.random.SeedGenerator`.

For me this implies that 

`keras.random.unform(shape=(), seed=None) == keras.random.unform(shape=(), seed=None)`

Which however is not the case.

See the colab here [https://colab.research.google.com/drive/19ITFpOVGBoVzBGWCXvMHsgUulVLjHuvm?usp=sharing](url) to see that calling keras.random.unform(shape=(), seed=None) produces a sequence of random numbers.

Am I misunderstanding this or is the documentation just wrong?


Pull request description:
Random seed doc
Following the discussion under #20569 this pull request contains a reformulation of the outdated doc concerning
the seed argument of all RNG functions. I also added a clarification to the SeedGenerator class documentation.

Resolves #20569 

Patch:
diff --git a/keras/src/backend/common/masking.py b/keras/src/backend/common/masking.py
index 63c5e85a0..afd0c2b64 100644
--- a/keras/src/backend/common/masking.py
+++ b/keras/src/backend/common/masking.py
@@ -3,8 +3,24 @@ from keras.src.backend.common.tensor_attributes import set_tensor_attr
 
 
 def set_keras_mask(x, mask):
-    return set_tensor_attr(x, "_keras_mask", mask)
+    """Sets the Keras mask attribute for the given tensor in-place.
+
+    Args:
+        x: Input tensor.
+        mask: The mask tensor to be set. If `None`, the `_keras_mask` attribute
+            will be cleared.
+    """
+    set_tensor_attr(x, "_keras_mask", mask)
 
 
 def get_keras_mask(x):
+    """Gets the Keras mask attribute from the given tensor.
+
+    Args:
+        x: Input tensor.
+
+    Returns:
+        The mask tensor associated with the input tensor, or `None` if no mask
+        has been set.
+    """
     return get_tensor_attr(x, "_keras_mask")
diff --git a/keras/src/layers/attention/multi_head_attention.py b/keras/src/layers/attention/multi_head_attention.py
index 60467149d..ad4d55d3a 100644
--- a/keras/src/layers/attention/multi_head_attention.py
+++ b/keras/src/layers/attention/multi_head_attention.py
@@ -528,6 +528,14 @@ class MultiHeadAttention(Layer):
         self._return_attention_scores = return_attention_scores
         if key is None:
             key = value
+
+        # Delete the masks because the masks are handled at the level of the
+        # layer
+        query_mask = backend.get_keras_mask(query)
+        backend.set_keras_mask(query, None)
+        backend.set_keras_mask(value, None)
+        backend.set_keras_mask(key, None)
+
         attention_mask = self._compute_attention_mask(
             query,
             value,
@@ -540,14 +548,14 @@ class MultiHeadAttention(Layer):
         #   N = `num_attention_heads`
         #   H = `size_per_head`
 
-        # `query` = [B, T, N ,H]
-        query = self._query_dense.call(query)
+        # `query` = [B, T, N, H]
+        query = self._query_dense(query)
 
         # `key` = [B, S, N, H]
-        key = self._key_dense.call(key)
+        key = self._key_dense(key)
 
         # `value` = [B, S, N, H]
-        value = self._value_dense.call(value)
+        value = self._value_dense(value)
         attention_output, attention_scores = self._compute_attention(
             query,
             key,
@@ -555,7 +563,11 @@ class MultiHeadAttention(Layer):
             attention_mask,
             training,
         )
-        attention_output = self._output_dense.call(attention_output)
+        attention_output = self._output_dense(attention_output)
+
+        # Set mask on output if needed
+        if query_mask is not None:
+            backend.set_keras_mask(attention_output, query_mask)
 
         if return_attention_scores:
             return attention_output, attention_scores
diff --git a/keras/src/layers/core/masking.py b/keras/src/layers/core/masking.py
index 2041fc82a..64483aefb 100644
--- a/keras/src/layers/core/masking.py
+++ b/keras/src/layers/core/masking.py
@@ -32,7 +32,7 @@ class Masking(Layer):
     inputs[:, 5, :] = 0.
 
     model = keras.models.Sequential()
-    model.add(keras.layers.Masking(mask_value=0.)
+    model.add(keras.layers.Masking(mask_value=0.0))
     model.add(keras.layers.LSTM(32))
     output = model(inputs)
     # The time step 3 and 5 will be skipped from LSTM calculation.
diff --git a/keras/src/random/random.py b/keras/src/random/random.py
index 705411bcb..33d5d593d 100644
--- a/keras/src/random/random.py
+++ b/keras/src/random/random.py
@@ -15,14 +15,19 @@ def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
             supported. If not specified, `keras.config.floatx()` is used,
             which defaults to `float32` unless you configured it otherwise (via
             `keras.config.set_floatx(float_dtype)`).
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     return backend.random.normal(
         shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed
@@ -51,14 +56,19 @@ def categorical(logits, num_samples, dtype="int32", seed=None):
             row of the input. This will be the second dimension of the output
             tensor's shape.
         dtype: Optional dtype of the output tensor.
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
 
     Returns:
         A 2-D tensor with (batch_size, num_samples).
@@ -94,14 +104,19 @@ def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
             supported. If not specified, `keras.config.floatx()` is used,
             which defaults to `float32` unless you configured it otherwise (via
             `keras.config.set_floatx(float_dtype)`)
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     if dtype and not backend.is_float_dtype(dtype):
         raise ValueError(
@@ -133,14 +148,19 @@ def randint(shape, minval, maxval, dtype="int32", seed=None):
             supported. If not specified, `keras.config.floatx()` is used,
             which defaults to `float32` unless you configured it otherwise (via
             `keras.config.set_floatx(float_dtype)`)
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     if dtype and not backend.is_int_dtype(dtype):
         raise ValueError(
@@ -169,14 +189,19 @@ def truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
             supported. If not specified, `keras.config.floatx()` is used,
             which defaults to `float32` unless you configured it otherwise (via
             `keras.config.set_floatx(float_dtype)`)
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     return backend.random.truncated_normal(
         shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed
@@ -198,14 +223,19 @@ def shuffle(x, axis=0, seed=None):
         x: The tensor to be shuffled.
         axis: An integer specifying the axis along which to shuffle. Defaults to
             `0`.
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     return backend.random.shuffle(x, axis=axis, seed=seed)
 
@@ -221,14 +251,19 @@ def gamma(shape, alpha, dtype=None, seed=None):
             supported. If not specified, `keras.config.floatx()` is used,
             which defaults to `float32` unless you configured it otherwise (via
             `keras.config.set_floatx(float_dtype)`).
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     return backend.random.gamma(shape, alpha=alpha, dtype=dtype, seed=seed)
 
@@ -251,14 +286,19 @@ def binomial(shape, counts, probabilities, dtype=None, seed=None):
             supported. If not specified, `keras.config.floatx()` is used,
             which defaults to `float32` unless you configured it otherwise (via
             `keras.config.set_floatx(float_dtype)`).
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     return backend.random.binomial(
         shape,
@@ -286,14 +326,19 @@ def beta(shape, alpha, beta, dtype=None, seed=None):
             supported. If not specified, `keras.config.floatx()` is used,
             which defaults to `float32` unless you configured it otherwise (via
             `keras.config.set_floatx(float_dtype)`).
-        seed: A Python integer or instance of
-            `keras.random.SeedGenerator`.
-            Used to make the behavior of the initializer
-            deterministic. Note that an initializer seeded with an integer
-            or None (unseeded) will produce the same random values
-            across multiple calls. To get different random values
-            across multiple calls, use as seed an instance
-            of `keras.random.SeedGenerator`.
+        seed: Optional Python integer or instance of
+           `keras.random.SeedGenerator`.
+            By default, the `seed` argument is `None`, and an internal global
+            `keras.random.SeedGenerator` is used. The `seed` argument can be
+            used to ensure deterministic (repeatable) random number generation.
+            Note that passing an integer as the `seed` value will produce the
+            same random values for each call. To generate different random
+            values for repeated calls, an instance of
+            `keras.random.SeedGenerator` must be provided as the `seed` value.
+            Remark concerning the JAX backend: When tracing functions with the
+            JAX backend the global `keras.random.SeedGenerator` is not
+            supported. Therefore, during tracing the default value seed=None
+            will produce an error, and a `seed` argument must be provided.
     """
     return backend.random.beta(
         shape=shape, alpha=alpha, beta=beta, dtype=dtype, seed=seed
diff --git a/keras/src/random/seed_generator.py b/keras/src/random/seed_generator.py
index 841f729f3..ad429e90b 100644
--- a/keras/src/random/seed_generator.py
+++ b/keras/src/random/seed_generator.py
@@ -11,14 +11,24 @@ from keras.src.utils.naming import auto_name
 
 @keras_export("keras.random.SeedGenerator")
 class SeedGenerator:
-    """Generates variable seeds upon each call to a RNG-using function.
-
-    In Keras, all RNG-using methods (such as `keras.random.normal()`)
-    are stateless, meaning that if you pass an integer seed to them
-    (such as `seed=42`), they will return the same values at each call.
-    In order to get different values at each call, you must use a
-    `SeedGenerator` instead as the seed argument. The `SeedGenerator`
-    object is stateful.
+    """Generates variable seeds upon each call to a function generating
+    random numbers.
+
+    In Keras, all random number generators (such as
+    `keras.random.normal()`) are stateless, meaning that if you pass an
+    integer seed to them (such as `seed=42`), they will return the same
+    values for repeated calls. To get different values for each
+    call, a `SeedGenerator` providing the state of the random generator
+    has to be used.
+    Note that all the random number generators have a default seed of None,
+    which implies that an internal global SeedGenerator is used.
+    If you need to decouple the RNG from the global state you can provide
+    a local `StateGenerator` with either a deterministic or random initial
+    state.
+    Remark concerning the JAX backen: Note that the use of a local
+    `StateGenerator` as seed argument is required for JIT compilation of
+    RNG with the JAX backend, because the use of global state is not
+    supported.
 
     Example:
 


Test patch:
diff --git a/keras/src/layers/attention/multi_head_attention_test.py b/keras/src/layers/attention/multi_head_attention_test.py
index 5cd374414..76e9cdbe6 100644
--- a/keras/src/layers/attention/multi_head_attention_test.py
+++ b/keras/src/layers/attention/multi_head_attention_test.py
@@ -1,4 +1,5 @@
 import os
+import warnings
 
 import numpy as np
 import pytest
@@ -315,6 +316,7 @@ class MultiHeadAttentionTest(testing.TestCase):
                 [[1, 2, 3, 0, 0], [3, 3, 1, 1, 2], [1, 0, 0, 0, 0]]
             )
             masked_query = layers.Embedding(4, 8, mask_zero=True)(query)
+            query_mask = backend.get_keras_mask(masked_query)
             value = np.random.normal(size=(3, 3, 8))
             output = layer(query=masked_query, value=value)
         except RuntimeError as e:
@@ -325,7 +327,7 @@ class MultiHeadAttentionTest(testing.TestCase):
                     "PyTorch errors out on GPU: issue to track bug is here "
                     "https://github.com/keras-team/keras/issues/20459"
                 )
-        self.assertAllClose(masked_query._keras_mask, output._keras_mask)
+        self.assertAllClose(query_mask, output._keras_mask)
 
     @parameterized.named_parameters(("causal", True), ("not_causal", 0))
     @pytest.mark.skipif(
@@ -373,6 +375,21 @@ class MultiHeadAttentionTest(testing.TestCase):
         self.assertAllClose(output_1, output_2)
         self.assertAllClose(output_1, output_3)
 
+    @pytest.mark.skipif(
+        backend.backend() == "numpy",
+        reason="Numpy backend does not support masking.",
+    )
+    def test_no_warning_with_keras_mask(self):
+        layer = layers.MultiHeadAttention(num_heads=2, key_dim=2)
+        query = np.array([[1, 2, 3, 0, 0], [3, 3, 1, 1, 2], [1, 0, 0, 0, 0]])
+        masked_query = layers.Embedding(4, 8, mask_zero=True)(query)
+        value = np.array([[5, 4, 0], [3, 0, 0], [2, 1, 1]])
+        masked_value = layers.Embedding(6, 8, mask_zero=True)(value)
+
+        with warnings.catch_warnings(record=True) as warning_logs:
+            _ = layer(query=masked_query, value=masked_value)
+            self.assertLen(warning_logs, 0)
+
     @parameterized.named_parameters(
         ("disable_flash_attention", False), ("enable_flash_attention", True)
     )
@@ -646,3 +663,31 @@ class MultiHeadAttentionTest(testing.TestCase):
     def test_multi_head_attention_output_shape_error(self):
         with self.assertRaisesRegex(ValueError, r"Invalid `output_shape`"):
             layers.MultiHeadAttention(num_heads=2, key_dim=16, output_shape=8.0)
+
+    def test_quantize_int8(self):
+        query = np.array([[[1.0, 0.0], [0.0, 1.0]]])
+        key = np.array([[[0.0, 1.0], [1.0, 0.0]]])
+        value = np.array([[[1.0, 2.0], [3.0, 4.0]]])
+        layer = layers.MultiHeadAttention(
+            num_heads=3,
+            key_dim=8,
+            use_bias=False,
+        )
+        layer.build(query.shape, value.shape, key.shape)
+        output_float = layer(query, key, value)
+        for sublayer in layer._flatten_layers():
+            try:
+                sublayer.quantize("int8")
+            except:
+                pass
+
+        # Verify weights dtype
+        self.assertDType(layer._query_dense._kernel, "int8")
+        self.assertDType(layer._key_dense._kernel, "int8")
+        self.assertDType(layer._value_dense._kernel, "int8")
+        self.assertDType(layer._output_dense._kernel, "int8")
+
+        # Try eager call and verify output correctness
+        output_quantized = layer(query, key, value)
+        mse = ops.mean(ops.square(output_float - output_quantized))
+        self.assertLess(mse, 1e-3)  # A weak correctness test