Updated documentation regarding DistributedEmbedding batch size. (#138)

The batch size represented by the first dimension of `input_shape` and `output_shape` in `FeatureConfig` is the global batch size.

Author: hertschuh

keras_rs/src/layers/embedding/base_distributed_embedding.py

@@ -146,14 +146,14 @@ class DistributedEmbedding(keras.layers.Layer):
     feature1 = keras_rs.layers.FeatureConfig(
         name="feature1",
         table=table1,
-        input_shape=(PER_REPLICA_BATCH_SIZE,),
-        output_shape=(PER_REPLICA_BATCH_SIZE, TABLE1_EMBEDDING_SIZE),
+        input_shape=(GLOBAL_BATCH_SIZE,),
+        output_shape=(GLOBAL_BATCH_SIZE, TABLE1_EMBEDDING_SIZE),
     )
     feature2 = keras_rs.layers.FeatureConfig(
         name="feature2",
         table=table2,
-        input_shape=(PER_REPLICA_BATCH_SIZE,),
-        output_shape=(PER_REPLICA_BATCH_SIZE, TABLE2_EMBEDDING_SIZE),
+        input_shape=(GLOBAL_BATCH_SIZE,),
+        output_shape=(GLOBAL_BATCH_SIZE, TABLE2_EMBEDDING_SIZE),
     )
 
     feature_configs = {

keras_rs/src/layers/embedding/distributed_embedding_config.py

@@ -102,7 +102,10 @@ class FeatureConfig:
         input_shape: The input shape of the feature. The feature fed into the
             layer has to match the shape. Note that for ragged dimensions in the
             input, the dimension provided here presents the maximum value;
-            anything larger will be truncated.
+            anything larger will be truncated. Also note that the first
+            dimension represents the global batch size. For example, on TPU,
+            this represents the total number of samples that are dispatched to
+            all the TPUs connected to the current host.
         output_shape: The output shape of the feature activation. What is
             returned by the embedding layer has to match this shape.
     """