diff --git a/imperative/python/megengine/amp/autocast.py b/imperative/python/megengine/amp/autocast.py index 03163034ded97199848cefc4e2292f9069101ccc..eeb23620e143bc8087094139e5cba707a1306eab 100644 --- a/imperative/python/megengine/amp/autocast.py +++ b/imperative/python/megengine/amp/autocast.py @@ -15,14 +15,14 @@ class autocast: A class to control autocast mode for amp as a context manager or a decorator. :param enabled: Whether autocast mode is enabled. - :low_prec_dtype: Set amp autocast mode's lower precision dtype. It will change the - target dtype in tensor casting for better speed and memory. Default: float16. - :high_prec_dtype: Set amp autocast mode's higher precision dtype. It will change the - target dtype in tensor casting for better precision. Default: float32. + :param low_prec_dtype: Set amp autocast mode's lower precision dtype. It will change + the target dtype in tensor casting for better speed and memory. Default: float16. + :param high_prec_dtype: Set amp autocast mode's higher precision dtype. It will + change the target dtype in tensor casting for better precision. Default: float32. Examples: - ..code-block:: + .. code-block:: # used as decorator @autocast() diff --git a/imperative/python/megengine/core/tensor/amp.py b/imperative/python/megengine/core/tensor/amp.py index 4f8e49a36e39c3c399a0e7079b81c25527e28ecf..9d54927b59fe049f099f0bc173b1eba08ff56ac5 100644 --- a/imperative/python/megengine/core/tensor/amp.py +++ b/imperative/python/megengine/core/tensor/amp.py @@ -17,7 +17,7 @@ def enabled(mod): Examples: - ..code-block:: + .. code-block:: import megengine as mge mge.amp.enabled = True @@ -40,7 +40,7 @@ def high_prec_dtype(mod): Examples: - ..code-block:: + .. code-block:: import megengine as mge mge.amp.high_prec_dtype = "float32" @@ -63,7 +63,7 @@ def low_prec_dtype(mod): Examples: - ..code-block:: + .. code-block:: import megengine as mge mge.amp.low_prec_dtype = "float16" diff --git a/imperative/python/megengine/quantization/qconfig.py b/imperative/python/megengine/quantization/qconfig.py index 5050045c2f62f852fe48936e02c21aba28b09f4a..6d5c476a992b66807f2ba61084e4f2bde1207d60 100644 --- a/imperative/python/megengine/quantization/qconfig.py +++ b/imperative/python/megengine/quantization/qconfig.py @@ -28,7 +28,7 @@ class QConfig( ) ): r""" - A config class indicating how to do quantize toward :class:`~.QATModule`'s + A config class indicating how to do quantize toward :class:`~.QATModule` 's ``activation`` and ``weight``. See :meth:`~.QATModule.set_qconfig` for detail usage. :param weight_observer: interface to instantiate an :class:`~.Observer` indicating diff --git a/imperative/python/megengine/tensor.py b/imperative/python/megengine/tensor.py index 8b09a1d180552d21468009f8bc443c2bc93c747b..2476355745bb39e3006cbbcd3c0b1bc448ab6419 100644 --- a/imperative/python/megengine/tensor.py +++ b/imperative/python/megengine/tensor.py @@ -30,6 +30,7 @@ class Tensor(_Tensor, ArrayMethodMixin): A tensor object represents a multidimensional, homogeneous array of fixed-size items. :param data: The value of returned Tensor. + :type data: Tensor, :class:`~.numpy.ndarray`, :class:`list` or python number. :param dtype: The dtype of returned Tensor. Uses data's dtype if not specified. :param device: The desired device of returned Tensor. Uses :func:`get_default_device` if not specified. :param is_const: Whether make it a ``ImutableTensor`` in tracing mode. @@ -43,7 +44,7 @@ class Tensor(_Tensor, ArrayMethodMixin): def __new__( cls, - data: Union["Tensor", np.ndarray, list, "scalar"] = None, + data: Union["Tensor", np.ndarray, list, int, float] = None, dtype: np.dtype = None, device: str = None, is_const: bool = False, @@ -76,7 +77,7 @@ class Tensor(_Tensor, ArrayMethodMixin): def __init__( self, - data: Union["Tensor", np.ndarray, list, "scalar"], + data: Union["Tensor", np.ndarray, list, int, float], dtype: np.dtype = None, device: str = None, is_const: bool = False,