Add docstrings in numpy (#617)

* Add docstrings in numpy

* passed review comments

Author: guillaumebaquiast

keras_core/ops/numpy.py

@@ -586,6 +586,38 @@ def append(
     x2,
     axis=None,
 ):
+    """Append tensor `x2` to the end of tensor `x1`.
+
+    Args:
+        x1: First input tensor.
+        x2: Second input tensor.
+        axis: Axis along which tensor `x2` is appended to tensor `x1`.
+            If `None`, both tensors are flattened before use.
+
+    Returns:
+        A tensor with the values of `x2` appended to `x1`.
+
+    Examples:
+
+    >>> x1 = keras_core.ops.convert_to_tensor([1, 2, 3])
+    >>> x2 = keras_core.ops.convert_to_tensor([[4, 5, 6], [7, 8, 9]])
+    >>> keras_core.ops.append(x1, x2)
+    array([1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)
+
+    When `axis` is specified, `x1` and `x2` must have compatible shapes.
+    >>> x1 = keras_core.ops.convert_to_tensor([[1, 2, 3], [4, 5, 6]])
+    >>> x2 = keras_core.ops.convert_to_tensor([[7, 8, 9]])
+    >>> keras_core.ops.append(x1, x2, axis=0)
+    array([[1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9]], dtype=int32)
+    >>> x3 = keras_core.ops.convert_to_tensor([7, 8, 9])
+    >>> keras_core.ops.append(x1, x3, axis=0)
+    Traceback (most recent call last):
+        ...
+    TypeError: Cannot concatenate arrays with different numbers of
+    dimensions: got (2, 3), (3,).
+    """
     if any_symbolic_tensors((x1, x2)):
         return Append(axis=axis).symbolic_call(x1, x2)
     return backend.numpy.append(x1, x2, axis=axis)
@@ -604,6 +636,51 @@ def compute_output_spec(self, start, stop=None, step=1, dtype=None):
 
 @keras_core_export(["keras_core.ops.arange", "keras_core.ops.numpy.arange"])
 def arange(start, stop=None, step=1, dtype=None):
+    """Return evenly spaced values within a given interval.
+
+    `arange` can be called with a varying number of positional arguments:
+    * `arange(stop)`: Values are generated within the half-open interval
+        `[0, stop)` (in other words, the interval including start but excluding
+        stop).
+    * `arange(start, stop)`: Values are generated within the half-open interval
+        `[start, stop)`.
+    * `arange(start, stop, step)`: Values are generated within the half-open
+        interval `[start, stop)`, with spacing between values given by step.
+
+    Args:
+        start: Integer or real, representing the start of the interval. The
+            interval includes this value.
+        stop: Integer or real, representing the end of the interval. The
+            interval does not include this value, except in some cases where
+            `step` is not an integer and floating point round-off affects the
+            lenght of `out`. Defaults to `None`.
+        step: Integer or real, represent the spacing between values. For any
+            output `out`, this is the distance between two adjacent values,
+            `out[i+1] - out[i]`. The default step size is 1. If `step` is
+            specified as a position argument, `start` must also be given.
+        dtype: The type of the output array. If `dtype` is not given, infer the
+            data type from the other input arguments.
+
+    Returns:
+        Tensor of evenly spaced values.
+        For floating point arguments, the length of the result is
+        `ceil((stop - start)/step)`. Because of floating point overflow, this
+        rule may result in the last element of out being greater than stop.
+
+    Examples:
+
+    >>> keras_core.ops.arange(3)
+    array([0, 1, 2], dtype=int32)
+
+    >>> keras_core.ops.arange(3.0)
+    array([0., 1., 2.], dtype=float32)
+
+    >>> keras_core.ops.arange(3, 7)
+    array([3, 4, 5, 6], dtype=int32)
+
+    >>> keras_core.ops.arange(3, 7, 2)
+    array([3, 5], dtype=int32)
+    """
     return backend.numpy.arange(start, stop, step=step, dtype=dtype)
 
 
@@ -617,6 +694,20 @@ def compute_output_spec(self, x):
 
 @keras_core_export(["keras_core.ops.arccos", "keras_core.ops.numpy.arccos"])
 def arccos(x):
+    """Trigonometric inverse cosine, element-wise. The inverse of `cos` so that,
+    if `y = cos(x)`, then `x = arccos(y)`.
+
+    Arg:
+        x: Input tensor.
+
+    Returns:
+        Tensor of the angle of the ray intersecting the unit circle at the given
+        x-coordinate in radians [0, pi].
+
+    Example:
+    >>> x = keras_core.ops.convert_to_tensor([1, -1])
+    array([0., 3.1415927], dtype=float32)
+    """
     if any_symbolic_tensors((x,)):
         return Arccos().symbolic_call(x)
     return backend.numpy.arccos(x)
@@ -632,6 +723,20 @@ def compute_output_spec(self, x):
 
 @keras_core_export(["keras_core.ops.arcsin", "keras_core.ops.numpy.arcsin"])
 def arcsin(x):
+    """Inverse sine, element-wise.
+
+    Arg:
+        x: Input tensor.
+
+    Returns:
+        Tensor of the inverse sine of each element in `x`, in radians and in
+        the closed interval `[-pi/2, pi/2]`.
+
+    Example:
+    >>> x = keras_core.ops.convert_to_tensor([1, -1, 0])
+    >>> keras_core.ops.arcsin(x)
+    array([ 1.5707964, -1.5707964,  0.], dtype=float32)
+    """
     if any_symbolic_tensors((x,)):
         return Arcsin().symbolic_call(x)
     return backend.numpy.arcsin(x)
@@ -647,6 +752,20 @@ def compute_output_spec(self, x):
 
 @keras_core_export(["keras_core.ops.arctan", "keras_core.ops.numpy.arctan"])
 def arctan(x):
+    """Trigonometric inverse tangent, element-wise.
+
+    Arg:
+        x: Input tensor.
+
+    Returns:
+        Tensor of the inverse tangent of each element in `x`, in the interval
+        `[-pi/2, pi/2]`.
+
+    Example:
+    >>> x = keras_core.ops.convert_to_tensor([0, 1])
+    >>> keras_core.ops.arctan(x)
+    array([0., 0.7853982], dtype=float32)
+    """
     if any_symbolic_tensors((x,)):
         return Arctan().symbolic_call(x)
     return backend.numpy.arctan(x)
@@ -665,6 +784,43 @@ def compute_output_spec(self, x1, x2):
 
 @keras_core_export(["keras_core.ops.arctan2", "keras_core.ops.numpy.arctan2"])
 def arctan2(x1, x2):
+    """Element-wise arc tangent of `x1/x2` choosing the quadrant correctly.
+
+    The quadrant (i.e., branch) is chosen so that `arctan2(x1, x2)` is the
+    signed angle in radians between the ray ending at the origin and passing
+    through the point `(1, 0)`, and the ray ending at the origin and passing
+    through the point `(x2, x1)`. (Note the role reversal: the "y-coordinate"
+    is the first function parameter, the "x-coordinate" is the second.) By IEEE
+    convention, this function is defined for `x2 = +/-0` and for either or both
+    of `x1` and `x2` `= +/-inf`.
+
+    Args:
+        x1: First input tensor.
+        x2: Second input tensor.
+
+    Returns:
+        Tensor of angles in radians, in the range `[-pi, pi]`.
+
+    Examples:
+    Consider four points in different quadrants:
+    >>> x = keras_core.ops.convert_to_tensor([-1, +1, +1, -1])
+    >>> y = keras_core.ops.convert_to_tensor([-1, -1, +1, +1])
+    >>> keras_core.ops.arctan2(y, x) * 180 / numpy.pi
+    array([-135., -45., 45., 135.], dtype=float32)
+
+    Note the order of the parameters. `arctan2` is defined also when x2=0 and
+    at several other points, obtaining values in the range `[-pi, pi]`:
+    >>> keras_core.ops.arctan2(
+    ...     keras_core.ops.array([1., -1.]),
+    ...     keras_core.ops.array([0., 0.]),
+    ... )
+    array([ 1.5707964, -1.5707964], dtype=float32)
+    >>> keras_core.ops.arctan2(
+    ...     keras_core.ops.array([0., 0., numpy.inf]),
+    ...     keras_core.ops.array([+0., -0., numpy.inf]),
+    ... )
+    array([0., 3.1415925, 0.7853982], dtype=float32)
+    """
     if any_symbolic_tensors((x1, x2)):
         return Arctan2().symbolic_call(x1, x2)
     return backend.numpy.arctan2(x1, x2)
@@ -688,6 +844,29 @@ def compute_output_spec(self, x):
 
 @keras_core_export(["keras_core.ops.argmax", "keras_core.ops.numpy.argmax"])
 def argmax(x, axis=None):
+    """Returns the indices of the maximum values along an axis.
+
+    Args:
+        x: Input tensor.
+        axis: By default, the index is into the flattened tensor, otherwise
+            along the specified axis.
+
+    Returns:
+        Tensor of indices. It has the same shape as `x`, with the dimension
+        along `axis` removed.
+
+    Example:
+    >>> x = keras_core.ops.arange(6).reshape(2, 3) + 10
+    >>> x
+    array([[10, 11, 12],
+           [13, 14, 15]], dtype=int32)
+    >>> keras_core.ops.argmax(x)
+    array(5, dtype=int32)
+    >>> keras_core.ops.argmax(x, axis=0)
+    array([1, 1, 1], dtype=int32)
+    >>> keras_core.ops.argmax(x, axis=1)
+    array([2, 2], dtype=int32)
+    """
     if any_symbolic_tensors((x,)):
         return Argmax(axis=axis).symbolic_call(x)
     return backend.numpy.argmax(x, axis=axis)
@@ -711,6 +890,29 @@ def compute_output_spec(self, x):
 
 @keras_core_export(["keras_core.ops.argmin", "keras_core.ops.numpy.argmin"])
 def argmin(x, axis=None):
+    """Returns the indices of the minium values along an axis.
+
+    Args:
+        x: Input tensor.
+        axis: By default, the index is into the flattened tensor, otherwise
+            along the specified axis.
+
+    Returns:
+        Tensor of indices. It has the same shape as `x`, with the dimension
+        along `axis` removed.
+
+    Example:
+    >>> x = keras_core.ops.arange(6).reshape(2, 3) + 10
+    >>> x
+    array([[10, 11, 12],
+           [13, 14, 15]], dtype=int32)
+    >>> keras_core.ops.argmin(x)
+    array(0, dtype=int32)
+    >>> keras_core.ops.argmin(x, axis=0)
+    array([0, 0, 0], dtype=int32)
+    >>> keras_core.ops.argmin(x, axis=1)
+    array([0, 0], dtype=int32)
+    """
     if any_symbolic_tensors((x,)):
         return Argmin(axis=axis).symbolic_call(x)
     return backend.numpy.argmin(x, axis=axis)
@@ -732,6 +934,37 @@ def compute_output_spec(self, x):
 
 @keras_core_export(["keras_core.ops.argsort", "keras_core.ops.numpy.argsort"])
 def argsort(x, axis=-1):
+    """Returns the indices that would sort an tensor.
+
+    Args:
+        x: Input tensor.
+        axis: Axis along which to sort. Default is `-1` (the last axis). If
+        `None`, the flattened tensor is used.
+
+    Returns:
+        Tensor of indices that sort `x` along the specified `axis`.
+
+    Examples:
+    One dimensional array:
+    >>> x = keras_core.ops.array([3, 1, 2])
+    >>> keras_core.ops.argsort(x)
+    array([1, 2, 0], dtype=int32)
+
+    Two-dimensional array:
+    >>> x = keras_core.ops.array([[0, 3], [3, 2], [4, 5]])
+    >>> x
+    array([[0, 3],
+           [3, 2],
+           [4, 5]], dtype=int32)
+    >>> keras_core.ops.argsort(x, axis=0)
+    array([[0, 1],
+           [1, 0],
+           [2, 2]], dtype=int32)
+    >>> keras_core.ops.argsort(x, axis=1)
+    array([[0, 1],
+           [1, 0],
+           [0, 1]], dtype=int32)
+    """
     if any_symbolic_tensors((x,)):
         return Argsort(axis=axis).symbolic_call(x)
     return backend.numpy.argsort(x, axis=axis)
@@ -747,6 +980,22 @@ def compute_output_spec(self, x, dtype=None):
 
 @keras_core_export(["keras_core.ops.array", "keras_core.ops.numpy.array"])
 def array(x, dtype=None):
+    """Create a tensor.
+
+    Args:
+        x: Input tensor.
+        dtype: The desired data-type for the tensor.
+
+    Returns:
+        A tensor.
+
+    Examples:
+    >>> keras_core.ops.array([1, 2, 3])
+    array([1, 2, 3], dtype=int32)
+
+    >>> keras_core.ops.array([1, 2, 3], dtype="float32")
+    array([1., 2., 3.], dtype=float32)
+    """
     if any_symbolic_tensors((x,)):
         return Array().symbolic_call(x, dtype=dtype)
     return backend.numpy.array(x, dtype=dtype)
@@ -800,6 +1049,57 @@ def compute_output_spec(self, x, weights=None):
 
 @keras_core_export(["keras_core.ops.average", "keras_core.ops.numpy.average"])
 def average(x, axis=None, weights=None):
+    """Compute the weighted average along the specified axis.
+
+    Args:
+        x: Input tensor.
+        axis: Integer along which to average `x`. The default, `axis=None`,
+            will average over all of the elements of the input tendor. If axis
+            is negative it counts from the last to the first axis.
+        weights: Tensor of wieghts associated with the values in `x`. Each
+            value in `x` contributes to the average according to its
+            associated weight. The weights array can either be 1-D (in which
+            case its length must be the size of a along the given axis) or of
+            the same shape as `x`. If `weights=None` (default), then all data
+            in `x` are assumed to have a weight equal to one.
+
+            The 1-D calculation is: `avg = sum(a * weights) / sum(weights)`.
+            The only constraint on weights is that `sum(weights)` must not be 0.
+
+    Returns:
+        Return the average along the specified axis.
+
+    Examples:
+    >>> data = keras_core.ops.arange(1, 5)
+    >>> data
+    array([1, 2, 3, 4], dtype=int32)
+    >>> keras_core.ops.average(data)
+    array(2.5, dtype=float32)
+    >>> keras_core.ops.average(
+    ...     keras_core.ops.arange(1, 11),
+    ...     weights=keras_core.ops.arange(10, 0, -1)
+    ... )
+    array(4., dtype=float32)
+
+    >>> data = keras_core.ops.arange(6).reshape((3, 2))
+    >>> data
+    array([[0, 1],
+           [2, 3],
+           [4, 5]], dtype=int32)
+    >>> keras_core.ops.average(
+    ...     data,
+    ...     axis=1,
+    ...     weights=keras_core.ops.array([1./4, 3./4])
+    ... )
+    array([0.75, 2.75, 4.75], dtype=float32)
+    >>> keras_core.ops.average(
+    ...     data,
+    ...     weights=keras_core.ops.array([1./4, 3./4])
+    ... )
+    Traceback (most recent call last):
+        ...
+    ValueError: Axis must be specified when shapes of a and weights differ.
+    """
     if any_symbolic_tensors((x,)):
         return Average(axis=axis).symbolic_call(x, weights=weights)
     return backend.numpy.average(x, weights=weights, axis=axis)