Enhanced Docstrings and Examples in /ops/function.py, /ops/math.py and /ops/nn.py.

keras_core/ops/function.py

@@ -180,7 +180,7 @@ def map_graph(inputs, outputs):
 
     Returns:
         A tuple `(nodes, nodes_by_depth, operations, operations_by_depth)`.
-        - nodes: list of Node instances.
+        - network_nodes: dict mapping unique node keys to the Node instances
         - nodes_by_depth: dict mapping ints (depth) to lists of node instances.
         - operations: list of Operation instances.
         - operations_by_depth: dict mapping ints (depth) to lists of Operation

keras_core/ops/math.py

@@ -48,10 +48,11 @@ def segment_sum(data, segment_ids, num_segments=None, sorted=False):
 
     Example:
 
-    >>> data = keras_core.ops.convert_to_tensor([1, 2, 3, 4, 5, 6])
-    >>> segment_ids = keras_core.ops.convert_to_tensor([0, 1, 0, 1, 0, 1])
-    >>> segment_sum(data, segment_ids)
-    array([9 12], shape=(2,), dtype=int32)
+    >>> data = keras_core.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
+    >>> segment_ids = keras_core.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
+    >>> num_segments = 3
+    >>> keras_core.ops.segment_sum(data, segment_ids,num_segments)
+    array([3, 30, 300], dtype=int32)
     """
     if any_symbolic_tensors((data,)):
         return SegmentSum(num_segments, sorted).symbolic_call(data, segment_ids)
@@ -100,10 +101,11 @@ def segment_max(data, segment_ids, num_segments=None, sorted=False):
 
     Example:
 
-    >>> data = keras_core.ops.convert_to_tensor([1, 2, 3, 4, 5, 6])
-    >>> segment_ids = keras_core.ops.convert_to_tensor([0, 1, 0, 1, 0, 1])
-    >>> segment_max(data, segment_ids)
-    array([9 12], shape=(2,), dtype=int32)
+    >>> data = keras_core.ops.convert_to_tensor([1, 2, 10, 20, 100, 200])
+    >>> segment_ids = keras_core.ops.convert_to_tensor([0, 0, 1, 1, 2, 2])
+    >>> num_segments = 3
+    >>> keras_core.ops.segment_max(data, segment_ids, num_segments)
+    array([2, 20, 200], dtype=int32)
     """
     if any_symbolic_tensors((data,)):
         return SegmentMax(num_segments, sorted).symbolic_call(data, segment_ids)
@@ -696,8 +698,9 @@ class Rsqrt(Operation):
 
     Example:
 
-    >>> x = keras_core.ops.convert_to_tensor([2., 3., -2.])
-    >>> rsqrt(x)
+    >>> data = keras_core.ops.convert_to_tensor([1.0, 10.0, 100.0])
+    >>> keras_core.ops.rsqrt(data)
+    array([1.0, 0.31622776, 0.1], dtype=float32)
     """
 
     def call(self, x):

keras_core/ops/nn.py

@@ -1501,6 +1501,30 @@ def compute_output_spec(self, inputs):
     ]
 )
 def multi_hot(inputs, num_tokens, axis=-1, dtype=None):
+    """Encodes integer labels as multi-hot vectors.
+
+    This function encodes integer labels as multi-hot vectors, where each label
+    is mapped to a binary value in the resulting vector.
+
+    Args:
+        inputs: Tensor of integer labels to be converted to multi-hot vectors.
+        num_tokens: Integer, the total number of unique tokens or classes.
+        axis: (optional) Axis along which the multi-hot encoding should be
+            added. Default is -1, which corresponds to the last dimension.
+        dtype: (optional) The data type of the resulting tensor. Default
+            is backend's float type.
+
+    Returns:
+        Tensor: The multi-hot encoded tensor.
+
+
+    Example:
+
+    >>> data = keras_core.ops.convert_to_tensor([0, 4])
+    >>> keras_core.ops.multi_hot(data, num_tokens=5)
+    array([1.0, 0.0, 0.0, 0.0, 1.0], dtype=float32)
+
+    """
     if any_symbolic_tensors((inputs,)):
         return MultiHot(num_tokens, axis, dtype).symbolic_call(inputs)