mindspore 2.0.0a0__cp39-cp39-win_amd64.whl → 2.0.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/function/grad/grad_func.py

@@ -17,7 +17,7 @@
 from __future__ import absolute_import
 from functools import partial
 import numpy as np
-from mindspore.common import jit
+from mindspore.common import jit, mutable
 from mindspore.common import Tensor
 from mindspore.common import dtype as mstype
 from mindspore.nn.cell import Cell
@@ -26,6 +26,7 @@ from mindspore.ops.primitive import constexpr
 from mindspore.ops.function.array_func import ones, expand_dims, size, reshape, broadcast_to, transpose
 from mindspore.ops.composite import _Vmap, _Grad, _TaylorOperation, GradOperation
 from mindspore.ops import operations as P
+from mindspore.ops.operations import _inner_ops as inner
 
 cast = P.Cast()
 dtype = P.DType()
@@ -86,11 +87,12 @@ def _check_grad_position(grad_position, args_num):
 
 
 @constexpr
-def _get_grad_op(get_by_list, get_by_position, has_aux, get_value=False):
-    return _Grad(get_by_list=get_by_list, get_by_position=get_by_position, has_aux=has_aux, get_value=get_value)
+def _get_grad_op(get_by_list, get_by_position, has_aux, get_value=False, return_ids=False):
+    return _Grad(get_by_list=get_by_list, get_by_position=get_by_position, has_aux=has_aux, get_value=get_value,
+                 return_ids=return_ids)
 
 
-def grad(fn, grad_position=0, weights=None, has_aux=False):
+def grad(fn, grad_position=0, weights=None, has_aux=False, return_ids=False):
     """
     A wrapper function to generate the gradient function for the input function.
 
@@ -113,11 +115,19 @@ def grad(fn, grad_position=0, weights=None, has_aux=False):
         has_aux (bool): If True, only the first output of `fn` contributes the gradient of `fn`, while the other outputs
             will be returned straightly. It means the `fn` must return more than one outputs in this case.
             Default: False.
+        return_ids(bool): Whether return the tuple made by gradients and the index to specify which inputs
+            to be differentiated or the name of parameters of the training network that need to calculate the gradient.
+            If True, the output gradients will be replaced by the tuples made by gradients and the index to specify
+            which inputs to be differentiated or the name of parameters of the training network.
+            Default: False.
 
     Returns:
         Function, the gradient function to calculate gradient for the input function or cell.
         For example, as for `out1, out2 = fn(*args)`, when `has_aux` is set True, gradient function will return outputs
         like `(gradient, out2)` and `out2` does not contribute to the differentiation, otherwise `gradient`.
+        When return_ids is set to True, The format of the output will be the same with the output of grad when
+        return_ids is set to false, but every gradient in the output will be replaced by a tuple of position id or
+        parameter name and its gradient.
 
     Raises:
         ValueError: If both `grad_position` and `weights` are None.
@@ -160,7 +170,7 @@ def grad(fn, grad_position=0, weights=None, has_aux=False):
         >>> print(aux)
         (Tensor(shape=[2], dtype=Float32, value= [ 5.00000000e+00,  5.00000000e+00]),)
         >>>
-        >>> # For given network to be differentiated with both inputs and weights, there are 3 cases.
+        >>> # For given network to be differentiated with both inputs and weights, there are 4 cases.
         >>> net = nn.Dense(10, 1)
         >>> loss_fn = nn.MSELoss()
         >>> def forward(inputs, labels):
@@ -192,17 +202,36 @@ def grad(fn, grad_position=0, weights=None, has_aux=False):
         >>> inputs_gradient, params_gradient = grad_fn(inputs, labels)
         >>> print(len(weights), len(params_gradient))
         2 2
+        >>> # Case 4: return the gradient with ids.
+        >>> import numpy as np
+        >>> import mindspore
+        >>> import mindspore.nn as nn
+        >>> from mindspore import Tensor, ops
+        >>> from mindspore import grad
+        >>>
+        >>> # Cell object to be differentiated
+        >>> class Net(nn.Cell):
+        ...     def construct(self, x, y, z):
+        ...         return x * y * z
+        >>> x = Tensor([1, 2], mindspore.float32)
+        >>> y = Tensor([-2, 3], mindspore.float32)
+        >>> z = Tensor([0, 3], mindspore.float32)
+        >>> net = Net()
+        >>> output = grad(net, grad_position=(1, 2), return_ids = True)(x, y, z)
+        >>> print(output)
+        ((1, Tensor(shape=[2], dtype=Float32, value=[ 0.00000000e+00,  6.00000000e+00])),
+         (2, Tensor(shape=[2], dtype=Float32, value=[-2.00000000e+00,  6.00000000e+00])))
     """
     if grad_position is None and weights is None:
         raise ValueError("`grad_position` and `weight` can not be None at the same time.")
 
     if grad_position is None:
-        return _get_grad_op(True, False, has_aux)(fn, weights)
+        return _get_grad_op(True, False, has_aux, False, return_ids)(fn, weights)
 
     grad_position = _convert_grad_position_type(grad_position)
     if weights is None:
-        return _get_grad_op(False, True, has_aux)(fn, None, grad_position)
-    return _get_grad_op(True, True, has_aux)(fn, weights, grad_position)
+        return _get_grad_op(False, True, has_aux, False, return_ids)(fn, None, grad_position)
+    return _get_grad_op(True, True, has_aux, False, return_ids)(fn, weights, grad_position)
 
 
 def value_and_grad(fn, grad_position=0, weights=None, has_aux=False):
@@ -329,6 +358,55 @@ def value_and_grad(fn, grad_position=0, weights=None, has_aux=False):
     return _get_grad_op(True, True, has_aux, True)(fn, weights, grad_position)
 
 
+def get_grad(gradients, identifier):
+    """
+    When `return_ids` of :func:`mindspore.grad` is set to True, use its return value as gradients. Then find
+    the specific gradient from `gradients` according to `identifier` .
+
+    As for gradient, two typical cases are included:
+
+    1. `identifier` is the position of the specific tensor to get gradient.
+    2. `identifier` is a parameter of a network.
+
+    Args:
+        gradients (Union[tuple[int, Tensor], tuple[tuple, tuple]]): The return value of :func:`mindspore.grad`
+            when `return_ids` is set to True.
+        identifier (Union[int, Parameter]): The position number of a tensor, or a parameter that is used in
+            :func:`mindspore.grad`.
+
+    Returns:
+        The gradient of the tensor on the position or in the parameter that specified by the `identifier`.
+
+    Raises:
+        RuntimeError: If gradient is not found.
+        TypeError: If type of Args does not belong to required ones.
+
+    Supported Platforms:
+        ``Ascend`` ``GPU`` ``CPU``
+
+    Examples:
+        >>> import numpy as np
+        >>> import mindspore
+        >>> import mindspore.nn as nn
+        >>> from mindspore import Tensor, ops
+        >>> from mindspore import grad, get_grad
+        >>>
+        >>>  # Cell object to be differentiated
+        >>> class Net(nn.Cell):
+        ...     def construct(self, x, y, z):
+        ...         return x * y * z
+        >>> x = Tensor([1, 2], mindspore.float32)
+        >>> y = Tensor([-2, 3], mindspore.float32)
+        >>> z = Tensor([0, 3], mindspore.float32)
+        >>> net = Net()
+        >>> out_grad = grad(net, grad_position=(1, 2), return_ids=True)(x, y, z)
+        >>> output = get_grad(out_grad, 1)
+        >>> print(output)
+        [0. 6.]
+    """
+    return inner.GetGrad()(gradients, identifier)
+
+
 def _trans_jet_inputs(primals_item, series_item):
     """Trans inputs of jet"""
     value_type = [mstype.int32, mstype.int64, mstype.float32, mstype.float64]
@@ -572,11 +650,16 @@ _grad_single = GradOperation(sens_param=True)
 _grad_all = GradOperation(sens_param=True, get_all=True)
 
 
+@constexpr
+def _check_jvp_input_v_len(inputs_len, v_len):
+    if inputs_len != v_len:
+        raise ValueError(f'v has invalid length: should be {inputs_len}, but got {v_len}')
+
+
 def jvp(fn, inputs, v, has_aux=False):
     """
     Compute the jacobian-vector-product of the given network. `jvp` matches
-    `forward-mode differentiation <https://www.mindspore.cn/docs/en/r2.0.0-alpha/design/auto_gradient.html
-    #forward-mode-ad>`_.
+    `forward-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#forward-mode-ad>`_.
 
     Args:
         fn (Union[Function, Cell]): The function or net that takes Tensor inputs and returns single Tensor or tuple of
@@ -662,9 +745,9 @@ def jvp(fn, inputs, v, has_aux=False):
         if isinstance(outputs, tuple):
             u = ()
             for item in outputs:
-                u = u + (oneslike(item),)
+                u = u + (mutable(oneslike(item)),)
         else:
-            u = oneslike(outputs)
+            u = mutable(oneslike(outputs))
         if len(jvp_inputs) == 1:
             second_grad_net = _grad_single(grad_single)
             gradient_outputs = second_grad_net(u, jvp_inputs, vectors)
@@ -689,6 +772,15 @@ def jvp(fn, inputs, v, has_aux=False):
 
     if not isinstance(inputs, (Tensor, tuple, list)) or not isinstance(v, (Tensor, tuple, list)):
         _raise_type_error()
+
+    inputs_len = 1
+    v_len = 1
+    if isinstance(inputs, (tuple, list)):
+        inputs_len = len(inputs)
+    if isinstance(v, (tuple, list)):
+        v_len = len(v)
+    _check_jvp_input_v_len(inputs_len, v_len)
+
     if isinstance(v, list):
         v = tuple(v)
     if isinstance(inputs, (tuple, list)):
@@ -776,8 +868,7 @@ def _check_tensor(inputs):
 def vjp(fn, *inputs, has_aux=False):
     """
     Compute the vector-jacobian-product of the given network. `vjp` matches
-    `reverse-mode differentiation <https://www.mindspore.cn/docs/en/r2.0.0-alpha/design/auto_gradient.html
-    #reverse-mode-ad>`_.
+    `reverse-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#reverse-mode-ad>`_.
 
     Args:
         fn (Union[Function, Cell]): The function or net that takes Tensor inputs and returns single Tensor or tuple of
@@ -889,15 +980,23 @@ def _jacfwd_trans_item(item, inputs_shape, grad_position):
     return output_wrt_input_all
 
 
-def _jacfwd_postprocess(x, inputs_shape, grad_position):
+def _jac_postprocess(x, shape, grad_position, mode):
     """reformat jacobian."""
+
+    if mode == 'forward':
+        func = _jacfwd_trans_item
+        args = (shape, grad_position)
+    else:
+        func = _jacrev_trans_item
+        args = (shape,)
+
     if isinstance(x, tuple):
         jacobian = ()
         for item in x:
-            jacobian += _jacfwd_trans_item(item, inputs_shape, grad_position)
+            jacobian += func(item, *args)
         res = jacobian
     else:
-        res = _jacfwd_trans_item(x, inputs_shape, grad_position)
+        res = func(x, *args)
     if len(res) == 1:
         return res[0]
     input_num = len(grad_position)
@@ -910,11 +1009,20 @@ def _jacfwd_postprocess(x, inputs_shape, grad_position):
     for i in range(output_num):
         input_grad = ()
         for j in range(input_num):
-            input_grad += (res[i * input_num + j],)
+            if mode == 'forward':
+                grad_increment = (res[i * input_num + j],)
+            else:
+                grad_increment = (res[j * output_num + i],)
+            input_grad += grad_increment
         jac += (input_grad,)
     return jac
 
 
+def _jacfwd_postprocess(x, inputs_shape, grad_position):
+    """reformat forward-computed Jacobian."""
+    return _jac_postprocess(x, inputs_shape, grad_position, 'forward')
+
+
 def _jacfwd_construct_v(inputs, grad_position):
     """
     For input (x1, x2), x1.shape = (a, b), x2.shape = (c, d), this method generates corresponding v (v1, v2),
@@ -927,10 +1035,9 @@ def _jacfwd_construct_v(inputs, grad_position):
     items_num = ()
     cum_num = (0,)
     for item in inputs:
-        item_num = size(item)
-        num += item_num
+        num += size(item)
         inputs_shape += ((item.shape, num),)
-        items_num += (item_num,)
+        items_num += (size(item),)
         cum_num += (num,)
     for i, element in enumerate(inputs):
         item_size = items_num[i]
@@ -953,18 +1060,17 @@ _vmap = _Vmap()
 def jacfwd(fn, grad_position=0, has_aux=False):
     """
     Compute Jacobian via forward mode, corresponding to
-    `forward-mode differentiation <https://www.mindspore.cn/docs/en/r2.0.0-alpha/design/auto_gradient.html
-    #forward-mode-ad>`_.
+    `forward-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#forward-mode-ad>`_.
     When number of outputs is much greater than that of inputs, it's better to calculate Jacobian via forward mode than
     reverse mode to get better performance.
 
     Args:
         fn (Union[Cell, Function]): Function to do GradOperation.
-        grad_position (Union[int, tuple[int]]): If int, get the gradient with respect to single input.
+        grad_position (Union[int, tuple[int]], optional): If int, get the gradient with respect to single input.
             If tuple, get the gradients with respect to selected inputs. 'grad_position' begins with 0. Default: 0.
-        has_aux (bool): If True, only the first output of `fn` contributes the gradient of `fn`, while the other outputs
-            will be returned straightly. It means the `fn` must return more than one outputs in this case.
-            Default: False.
+        has_aux (bool, optional): If True, only the first output of `fn` contributes the gradient of `fn`,
+            while the other outputs will be returned straightly. It means the `fn` must return more than one
+            outputs in this case. Default: False.
 
     Returns:
         Function, returns the Jacobian function for the input function or cell.
@@ -1032,9 +1138,9 @@ def jacfwd(fn, grad_position=0, has_aux=False):
             if isinstance(outputs, tuple):
                 u = ()
                 for item in outputs:
-                    u = u + (oneslike(item),)
+                    u = u + (mutable(oneslike(item)),)
             else:
-                u = oneslike(outputs)
+                u = mutable(oneslike(outputs))
             if len(jvp_inputs) == 1:
                 second_grad_net = _grad_single(grad_single)
             else:
@@ -1044,7 +1150,7 @@ def jacfwd(fn, grad_position=0, has_aux=False):
 
         def inner_aux_fn(jvp_inputs, vectors):
             outputs = aux_fn(*jvp_inputs)
-            u = oneslike(outputs)
+            u = mutable(oneslike(outputs))
             if len(jvp_inputs) == 1:
                 second_grad_net = _grad_single(grad_single)
             else:
@@ -1080,29 +1186,8 @@ def _jacrev_trans_item(item, outputs_shape):
 
 
 def _jacrev_postprocess(x, outputs_shape, grad_position):
-    """reformat jacobian."""
-    if isinstance(x, tuple):
-        jacobian = ()
-        for item in x:
-            jacobian += _jacrev_trans_item(item, outputs_shape)
-        res = jacobian
-    else:
-        res = _jacrev_trans_item(x, outputs_shape)
-    if len(res) == 1:
-        return res[0]
-    input_num = len(grad_position)
-    if len(res) % input_num != 0:
-        raise ValueError("The numbers of inputs and outputs do not match.")
-    output_num = len(res) // input_num
-    if input_num == 1 or output_num == 1:
-        return res
-    jac = ()
-    for i in range(output_num):
-        input_grad = ()
-        for j in range(input_num):
-            input_grad += (res[j * output_num + i],)
-        jac += (input_grad,)
-    return jac
+    """reformat reverse-computed jacobian."""
+    return _jac_postprocess(x, outputs_shape, grad_position, 'reverse')
 
 
 def _jacrev_construct_v(inputs, outputs, has_aux=False):
@@ -1126,7 +1211,7 @@ def _jacrev_construct_v(inputs, outputs, has_aux=False):
         outputs_shape += ((item.shape, num),)
         items_num += (item_num,)
         cum_num += (num,)
-    for i, element in enumerate(inputs):
+    for element in inputs:
         primal = broadcast_to(element, (num,) + element.shape)
         primals += (primal,)
     for i, element in enumerate(outputs):
@@ -1145,18 +1230,17 @@ _grad = _Grad(get_by_position=True, has_aux=False, sens_param=True)
 def jacrev(fn, grad_position=0, has_aux=False):
     """
     Compute Jacobian via reverse mode, corresponding to
-    `reverse-mode differentiation <https://www.mindspore.cn/docs/en/r2.0.0-alpha/design/auto_gradient.html
-    #reverse-mode-ad>`_.
+    `reverse-mode differentiation <https://www.mindspore.cn/docs/en/r2.0/design/auto_gradient.html#reverse-mode-ad>`_.
     When number of inputs is much greater than that of outputs, it's better to calculate Jacobian via reverse mode than
     forward mode to get better performance.
 
     Args:
         fn (Union[Cell, Function]): Function to do GradOperation.
-        grad_position (Union[int, tuple[int]]): If int, get the gradient with respect to single input.
+        grad_position (Union[int, tuple[int]], optional): If int, get the gradient with respect to single input.
             If tuple, get the gradients with respect to selected inputs. 'grad_position' begins with 0. Default: 0.
-        has_aux (bool): If True, only the first output of `fn` contributes the gradient of `fn`, while the other outputs
-            will be returned straightly. It means the `fn` must return more than one outputs in this case.
-            Default: False.
+        has_aux (bool, optional): If True, only the first output of `fn` contributes the gradient of `fn`,
+            while the other outputs will be returned straightly. It means the `fn` must return more than
+            one outputs in this case. Default: False.
 
     Returns:
         Function, returns the Jacobian function for the input function or cell.
@@ -1277,7 +1361,7 @@ def stop_gradient(value):
     StopGradient is used for eliminating the effect of a value on the gradient, such as truncating
     the gradient propagation from an output of a function.
     For more details, please refer to `Stop Gradient
-    <https://www.mindspore.cn/tutorials/en/r2.0.0-alpha/beginner/autograd.html#stop-gradient>`_.
+    <https://www.mindspore.cn/tutorials/en/r2.0/beginner/autograd.html#stop-gradient>`_.
 
     Args:
         value (Any): The value whose effect on the gradient to be eliminated.
@@ -1319,6 +1403,7 @@ __all__ = [
     'jvp',
     'vjp',
     'linearize',
-    'stop_gradient'
+    'stop_gradient',
+    'get_grad'
 ]
 __all__.sort()

mindspore/ops/function/image_func.py

@@ -26,18 +26,17 @@ from .._primitive_cache import _get_cache_prim
 def bounding_box_decode(anchor_box, deltas, max_shape, means=(0.0, 0.0, 0.0, 0.0), stds=(1.0, 1.0, 1.0, 1.0),
                         wh_ratio_clip=0.016):
     r"""
-    Decodes bounding boxes locations.
-
-    The function of the operator is to calculate the offset, and this operator converts the offset into a Bbox,
-    which is used to mark the target in the subsequent images, etc.
+    Decode the bounding box locations, calculate the offset, and convert the offset into a Bbox, which is used to mark
+    the target in the subsequent images, etc.
 
     Args:
         anchor_box (Tensor): Anchor boxes. The shape of `anchor_box` must be :math:`(n, 4)`.
         deltas (Tensor): Delta of boxes. Which has the same shape with `anchor_box`.
         max_shape (tuple): The max size limit for decoding box calculation.
-        means (tuple): The means of `deltas` calculation. Default: (0.0, 0.0, 0.0, 0.0).
-        stds (tuple): The standard deviations of `deltas` calculation. Default: (1.0, 1.0, 1.0, 1.0).
-        wh_ratio_clip (float): The limit of width and height ratio for decoding box calculation. Default: 0.016.
+        means (tuple, optional): The means of `deltas` calculation. Default: (0.0, 0.0, 0.0, 0.0).
+        stds (tuple, optional): The standard deviations of `deltas` calculation. Default: (1.0, 1.0, 1.0, 1.0).
+        wh_ratio_clip (float, optional): The limit of width and height ratio for decoding box calculation.
+            Default: 0.016.
 
     Returns:
         Tensor, decoded boxes. It has the same data type and shape as `anchor_box`.
@@ -66,16 +65,14 @@ def bounding_box_decode(anchor_box, deltas, max_shape, means=(0.0, 0.0, 0.0, 0.0
 
 def bounding_box_encode(anchor_box, groundtruth_box, means=(0.0, 0.0, 0.0, 0.0), stds=(1.0, 1.0, 1.0, 1.0)):
     r"""
-    Encodes bounding boxes locations.
-
-    This operator will calculate the offset between the predicted bounding boxes and the real bounding boxes,
-    and this offset will be used as a variable for the loss.
+    Encode the bounding box locations, calculate the offset between the predicted bounding boxes and
+    the real bounding boxes, and the offset will be used as a variable for the loss.
 
     Args:
         anchor_box (Tensor): Anchor boxes. The shape of `anchor_box` must be :math:`(n, 4)`.
         groundtruth_box (Tensor): Ground truth boxes. Which has the same shape with `anchor_box`.
-        means (tuple): Means for encoding bounding boxes calculation. Default: (0.0, 0.0, 0.0, 0.0).
-        stds (tuple): The standard deviations of deltas calculation. Default: (1.0, 1.0, 1.0, 1.0).
+        means (tuple, optional): Means for encoding bounding boxes calculation. Default: (0.0, 0.0, 0.0, 0.0).
+        stds (tuple, optional): The standard deviations of deltas calculation. Default: (1.0, 1.0, 1.0, 1.0).
 
     Returns:
         Tensor, encoded bounding boxes. It has the same data type and shape as input `anchor_box`.
@@ -103,18 +100,36 @@ def bounding_box_encode(anchor_box, groundtruth_box, means=(0.0, 0.0, 0.0, 0.0),
 
 def check_valid(bboxes, img_metas):
     r"""
-    Checks bounding box.
+    Checks whether the bounding box is in the image.
+
+    `bboxes` contain several sets of bounding boxes, each represented by two abscissa points :math:`(x0, x1)` and
+    two ordinate points :math:`(y0, y1)` .
+    `img_metas` provides information about the original image, including three parameters
+    :math:`(height, width, ratio)` , which specify the valid boundary of the image.
+
+    when the following conditions are met:
+
+    :math:`x0 >= 0`
 
-    Checks whether the bounding box cross data and data border are valid.
+    :math:`y0 >= 0`
+
+    :math:`x1 <= width * ratio - 1`
+
+    :math:`y1 <= height * ratio - 1`
+
+    the bounding box is considered to be within the image.
 
     .. warning::
-        Boundary(heights * ratio, widths * ratio) specified by `bboxes` is required to be valid.
+        The bounding box specified by `bboxes` and the image information specified by `img_metas` need to be valid,
+        i.e.:
+        :math:`x0 <= x1` , :math:`y0 <= y1` , and :math:`(height, width, ratio)` are all positive.
 
     Args:
-        bboxes (Tensor): Bounding boxes tensor with shape :math:`(N, 4)`. :math:`N` indicates the number of
-            bounding boxes, the value `4` indicates `x0`, `x1`, `y0`, and `y1`. Data type must be float16 or float32.
-        img_metas (Tensor): Raw image size information with the format of `(height, width, ratio)`, specifying
-            the valid boundary `(height * ratio, width * ratio)`. Data type must be float16 or float32.
+        bboxes (Tensor): Bounding boxes tensor with shape :math:`(N, 4)` . :math:`N` indicates the number of
+            bounding boxes, the value `4` indicates four coordinate points :math:`(x0, y0, x1, y1)` . Data type must
+            be float16 or float32.
+        img_metas (Tensor): Raw image size information with the format of :math:`(height, width, ratio)` , specifying
+            the valid boundary :math:`(height * ratio - 1, width * ratio - 1)` . Data type must be float16 or float32.
 
     Returns:
         Tensor, with shape of :math:`(N,)` and dtype of bool, specifying whether the bounding boxes is in the image.
@@ -139,7 +154,7 @@ def check_valid(bboxes, img_metas):
 
 
 def crop_and_resize(image, boxes, box_indices, crop_size, method="bilinear", extrapolation_value=0.0):
-    """
+    r"""
     Extracts crops from the input image Tensor and resizes them.
 
     Note:
@@ -147,28 +162,29 @@ def crop_and_resize(image, boxes, box_indices, crop_size, method="bilinear", ext
         For now, the backward of the operator only support bilinear method, for other methods, will return 0.
 
     Args:
-        image (Tensor): The input image must be a 4-D tensor of shape (batch, image_height, image_width, depth).
-           Types allowed: int8, int16, int32, int64, float16, float32, float64, uint8, uint16.
-        boxes (Tensor): A 2-D tensor of shape [num_boxes, 4].
-           The i-th row of the tensor specifies the coordinates of a box in the box_ind[i] image
-           and is specified in normalized coordinates [y1, x1, y2, x2]. A normalized coordinate value of y is mapped to
-           the image coordinate at y * (image_height - 1), so as the [0, 1] interval of normalized image height is
-           mapped to [0, image_height - 1] in image height coordinates. We do allow y1 > y2, in which case the sampled
-           crop is an up-down flipped version of the original image. The width dimension is treated similarly.
-           Normalized coordinates outside the [0, 1] range are allowed, in which case we use extrapolation_value to
-           extrapolate the input image values. Types allowed: float32.
-        box_indices (Tensor): A 1-D tensor of shape [num_boxes] with int32 values in [0, batch).
-           The value of box_ind[i] specifies the image that the i-th box refers to. Types allowed: int32.
-        crop_size (Tuple[int]): A tuple of two int32 elements: (crop_height, crop_width).
-           Only constant value is allowed. All cropped image patches are resized to this size.
-           The aspect ratio of the image content is not preserved. Both crop_height and crop_width need to be positive.
+        image (Tensor): A 4-D Tensor representing a batch of images. It has shape
+            :math:`(batch, image\_height, image\_width, depth)`.
+        boxes (Tensor):  A 2-D Tensor with shape :math:`(num\_boxes, 4)` representing the normalized
+            coordinates of the boxes to be cropped. The coordinates are specified in the
+            form :math:`[y1, x1, y2, x2]`, where :math:`(y1, x1)` is the first corner
+            and :math:`(y2, x2)` is the second corner of the box.
+            If :math:`y1 > y2`, the sampled crop is inverted upside down, the width dimensionis treated
+            similarly when :math:`x1 > x2`. If normalized coordinates are not in range :math:`[0, 1]`,
+            extrapolated input image values are used instead. Supported data type: float32.
+        box_indices (Tensor): A 1-D Tensor of shape :math:`\text{num\_boxes}` representing the batch
+            index for each box. Supported type: int32.
+        crop_size (Tuple[int]): A tuple of two elements: (crop_height, crop_width), representing
+            the output size of the cropped and resized images.
+            Only positive values are supported. Supported type: int32.
         method (str, optional): An optional string that specifies the sampling method for resizing.
            It can be "bilinear", "nearest" or "bilinear_v2". The option "bilinear" stands for standard bilinear
-           interpolation algorithm, while "bilinear_v2" may result in better result in some cases. Default: "bilinear"
+           interpolation algorithm, while "bilinear_v2" may result in better result in some cases.
+           "nearest" is the nearest neighbor interpolation algorithm.
+           Default: "bilinear".
         extrapolation_value (float, optional): An optional float value used extrapolation, if applicable. Default: 0.0.
 
     Returns:
-        A 4-D tensor of shape [num_boxes, crop_height, crop_width, depth] with type(float32).
+        A 4-D tensor of shape :math:`(num_boxes, crop_height, crop_width, depth)` with type(float32).
 
     Raises:
         TypeError: If `image` or `boxes` or `box_indices` is not a Tensor.