Refactor random namespace

keras_core/backend/common/__init__.py

@@ -1,8 +1,8 @@
 from keras_core.backend.common import backend_utils
-from keras_core.backend.common import random
 from keras_core.backend.common.variables import AutocastScope
 from keras_core.backend.common.variables import KerasVariable
 from keras_core.backend.common.variables import get_autocast_scope
 from keras_core.backend.common.variables import is_float_dtype
 from keras_core.backend.common.variables import standardize_dtype
 from keras_core.backend.common.variables import standardize_shape
+from keras_core.random import random

keras_core/backend/common/keras_tensor.py

@@ -4,6 +4,7 @@ from keras_core.api_export import keras_core_export
 from keras_core.utils.naming import auto_name
 
 
+@keras_core_export("keras_core.KerasTensor")
 class KerasTensor:
     def __init__(self, shape, dtype="float32", record_history=True, name=None):
         from keras_core import backend

keras_core/backend/common/stateless_scope.py

@@ -1,8 +1,10 @@
+from keras_core.api_export import keras_core_export
 from keras_core.backend.common import global_state
 from keras_core.backend.common.variables import KerasVariable
 from keras_core.backend.common.variables import initialize_all_variables
 
 
+@keras_core_export("keras_core.StatelessScope")
 class StatelessScope:
     def __init__(self, state_mapping=None, collect_losses=False):
         from keras_core import backend

keras_core/backend/common/variables.py

@@ -33,11 +33,18 @@ class KerasVariable:
                     "You are attempting to create a variable "
                     "while in a stateless scope. This is disallowed. "
                     "Make sure that all variables are created "
-                    "before you start using your layer/model objects. "
-                    "Most of this time, this means you need to "
+                    "before you start using your layer/model objects.\n\n"
+                    "In some cases, you might be seeing this error "
+                    "because you need to "
                     "implement a `def build(self, input_shape)` method "
                     "on your layer/model, which will "
-                    "create its variables."
+                    "create its variables.\n\n"
+                    "In some other cases, you might be seeing this error "
+                    "because you are instantiating a `Variable` and "
+                    "assigning it to a layer without going through "
+                    "self.add_variable()/self.add_weight(). Always prefer "
+                    "using these methods "
+                    "(with a `shape` and `initializer` argument)."
                 )
         else:
             if callable(initializer):

keras_core/backend/jax/random.py

@@ -1,9 +1,9 @@
 import jax
 
-from keras_core.backend.common.random import SeedGenerator
-from keras_core.backend.common.random import draw_seed
-from keras_core.backend.common.random import make_default_seed
 from keras_core.backend.config import floatx
+from keras_core.random.seed_generator import SeedGenerator
+from keras_core.random.seed_generator import draw_seed
+from keras_core.random.seed_generator import make_default_seed
 
 
 def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):

keras_core/backend/tensorflow/random.py

@@ -1,9 +1,9 @@
 import tensorflow as tf
 
-from keras_core.backend.common.random import SeedGenerator
-from keras_core.backend.common.random import draw_seed
-from keras_core.backend.common.random import make_default_seed
 from keras_core.backend.config import floatx
+from keras_core.random.seed_generator import SeedGenerator
+from keras_core.random.seed_generator import draw_seed
+from keras_core.random.seed_generator import make_default_seed
 
 
 def tf_draw_seed(seed):

keras_core/backend/torch/random.py

@@ -1,7 +1,7 @@
-from keras_core.backend.common.random import SeedGenerator
-from keras_core.backend.common.random import draw_seed
-from keras_core.backend.common.random import make_default_seed
 from keras_core.backend.config import floatx
+from keras_core.random.seed_generator import SeedGenerator
+from keras_core.random.seed_generator import draw_seed
+from keras_core.random.seed_generator import make_default_seed
 
 
 def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):

keras_core/operations/function.py

@@ -2,10 +2,12 @@ import collections
 
 from tensorflow import nest
 
+from keras_core.api_export import keras_core_export
 from keras_core.backend import KerasTensor
 from keras_core.operations.operation import Operation
 
 
+@keras_core_export("keras_core.Function")
 class Function(Operation):
     def __init__(self, inputs, outputs, name=None):
         super().__init__(name=name)

keras_core/operations/operation.py

@@ -4,6 +4,7 @@ import textwrap
 from tensorflow import nest
 
 from keras_core import backend
+from keras_core.api_export import keras_core_export
 from keras_core.backend.common.keras_tensor import any_symbolic_tensors
 from keras_core.operations.node import Node
 from keras_core.utils import python_utils
@@ -11,6 +12,7 @@ from keras_core.utils import traceback_utils
 from keras_core.utils.naming import auto_name
 
 
+@keras_core_export("keras_core.Operation")
 class Operation:
     def __init__(self, name=None):
         if name is None:

keras_core/operations/random.py

@@ -1,6 +0,0 @@
-from keras_core import backend
-
-dropout = backend.random.dropout
-normal = backend.random.normal
-truncated_normal = backend.random.truncated_normal
-uniform = backend.random.uniform

keras_core/random/random.py

@@ -0,0 +1,100 @@
+from keras_core import backend
+from keras_core.api_export import keras_core_export
+
+
+@keras_core_export("keras_core.random.normal")
+def normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
+    """Draw random samples from a normal (Gaussian) distribution.
+
+    Args:
+        shape: The shape of the random values to generate.
+        mean: Floats, defaults to 0. Mean of the random values to generate.
+        stddev: Floats, defaults to 1. Standard deviation of the random values
+            to generate.
+        dtype: Optional dtype of the tensor. Only floating point types are
+            supported. If not specified, `keras_core.config.floatx()` is used,
+            which defaults to `float32` unless you configured it otherwise (via
+            `keras_core.config.set_floatx(float_dtype)`).
+        seed: A Python integer or instance of
+            `keras_core.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_core.random.SeedGenerator`.
+    """
+    return normal(shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed)
+
+
+@keras_core_export("keras_core.random.uniform")
+def uniform(shape, minval=0.0, maxval=1.0, dtype=None, seed=None):
+    """Draw samples from a uniform distribution.
+
+    The generated values follow a uniform distribution in the range
+    `[minval, maxval)`. The lower bound `minval` is included in the range,
+    while the upper bound `maxval` is excluded.
+
+    For floats, the default range is `[0, 1)`.  For ints, at least `maxval`
+    must be specified explicitly.
+
+    Args:
+        shape: The shape of the random values to generate.
+        minval: Floats, defaults to 0. Lower bound of the range of
+            random values to generate (inclusive).
+        maxval: Floats, defaults to 1. Upper bound of the range of
+            random values to generate (exclusive).
+        dtype: Optional dtype of the tensor. Only floating point types are
+            supported. If not specified, `keras_core.config.floatx()` is used,
+            which defaults to `float32` unless you configured it otherwise (via
+            `keras_core.config.set_floatx(float_dtype)`)
+        seed: A Python integer or instance of
+            `keras_core.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_core.random.SeedGenerator`.
+    """
+    return backend.random.uniform(
+        shape, minval=minval, maxval=maxval, dtype=dtype, seed=seed
+    )
+
+
+@keras_core_export("keras_core.random.truncated_normal")
+def truncated_normal(shape, mean=0.0, stddev=1.0, dtype=None, seed=None):
+    """Draw samples from a truncated normal distribution.
+
+    The values are drawn from a normal distribution with specified mean and
+    standard deviation, discarding and re-drawing any samples that are more
+    than two standard deviations from the mean.
+
+    Args:
+        shape: The shape of the random values to generate.
+        mean: Floats, defaults to 0. Mean of the random values to generate.
+        stddev: Floats, defaults to 1. Standard deviation of the random values
+            to generate.
+        dtype: Optional dtype of the tensor. Only floating point types are
+            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_core.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_core.random.SeedGenerator`.
+    """
+    return backend.random.truncated_normal(
+        shape, mean=mean, stddev=stddev, dtype=dtype, seed=seed
+    )
+
+
+@keras_core_export("keras_core.random.dropout")
+def dropout(inputs, rate, noise_shape=None, seed=None):
+    return backend.random.dropout(
+        inputs, rate, noise_shape=noise_shape, seed=seed
+    )

keras_core/random/seed_generator.py

@@ -1,6 +1,9 @@
 import random as python_random
 
+from keras_core.api_export import keras_core_export
 
+
+@keras_core_export("keras_core.random.SeedGenerator")
 class SeedGenerator:
     def __init__(self, seed):
         from keras_core.backend import Variable