提交 d814f196 编写于 作者: 粉红谐星's avatar 粉红谐星

docs(mge/functional): update functional.zeros docstring

上级 719b5e69
......@@ -1587,9 +1587,7 @@ def one_hot(inp: Tensor, num_classes: int) -> Tensor:
[0 0 1 0]
[0 0 0 1]]
"""
zeros_tensor = zeros(
list(inp.shape) + [num_classes], dtype=inp.dtype, device=inp.device
)
zeros_tensor = zeros(list(inp.shape) + [num_classes], dtype=inp.dtype, device=inp.device)
ones_tensor = ones(list(inp.shape) + [1], dtype=inp.dtype, device=inp.device)
op = builtin.IndexingSetOneHot(axis=inp.ndim)
......
......@@ -195,14 +195,28 @@ def ones(
return full(shape, 1.0, dtype=dtype, device=device)
def zeros(shape, dtype="float32", device=None) -> Tensor:
r"""Returns a zero tensor with given shape.
def zeros(
shape: Union[int, Tuple[int, ...]],
*,
dtype="float32",
device: Optional[CompNode] = None
) -> Tensor:
r"""Returns a new tensor having a specified shape and filled with zeros.
Args:
shape: a list, tuple or integer defining the shape of the output tensor.
dtype: the desired data type of the output tensor. Default: ``float32``.
device: the desired device of the output tensor. Default: if ``None``,
use the default device (see :func:`~.megengine.get_default_device`).
shape (int or sequence of ints): the shape of the output tensor.
Keyword args:
dtype (:attr:`.Tensor.dtype`): output tensor data type. Default: ``float32``.
device (:attr:`.Tensor.device`): device on which to place the created tensor. Default: ``None``.
Returns:
a tensor containing zeros.
Examples:
>>> F.zeros((2, 1))
Tensor([[0.]
[0.]], device=xpux:0)
"""
return full(shape, 0.0, dtype=dtype, device=device)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册