mindspore 2.6.0rc1__cp39-cp39-win_amd64.whl → 2.7.0rc1__cp39-cp39-win_amd64.whl

mindspore/ops/composite/multitype_ops/add_impl.py

@@ -20,9 +20,19 @@ from mindspore.ops.composite.multitype_ops import _compile_utils as utils
 from mindspore.ops.composite import base
 from mindspore.ops import functional as F
 from mindspore.ops.composite.multitype_ops._constexpr_utils import make_tensor, check_equal
+from mindspore.ops.auto_generate.gen_ops_prim import InplaceAddExt, InplaceAddsExt
 from mindspore.common import CSRTensor, COOTensor
 from ...operations._sequence_ops import SequenceAdd
 
+# x += y
+augassign_add = base.MultitypeFuncGraph('augassign_add', True)
+"""
+`augassign_add` is a metafuncgraph object which will add two objects according to input type using ".register"
+decorator.
+"""
+
+augassign_add.set_need_raise()
+
 
 add = base.MultitypeFuncGraph('add', True)
 """`add` is a metafuncgraph object which will add two objects according to input type using ".register" decorator."""
@@ -103,7 +113,7 @@ class _ListAdd(base.ListAdd_):
 _list_add = _ListAdd('list_add')
 """`_list_add` is an metafuncgraph object which will concatenate two lists to form a list."""
 
-
+@augassign_add.register("Number", "Number")
 @add.register("Number", "Number")
 @_add_backward.register("Number", "Number")
 def _scalar_add_scalar(x, y):
@@ -119,7 +129,7 @@ def _scalar_add_scalar(x, y):
     """
     return F.scalar_add(x, y)
 
-
+@augassign_add.register("String", "String")
 @add.register("String", "String")
 def _string_concat_string(x, y):
     """
@@ -135,6 +145,7 @@ def _string_concat_string(x, y):
     return F.string_concat(x, y)
 
 
+@augassign_add.register("Number", "Tensor")
 @add.register("Number", "Tensor")
 def _scalar_add_tensor(x, y):
     """
@@ -150,6 +161,21 @@ def _scalar_add_tensor(x, y):
     return F.add(x, y)
 
 
+@augassign_add.register("Tensor", "Number")
+def _tensor_add_scalar_augassign(x, y):
+    """
+    Tensor is added to number.
+
+    Args:
+        x (Tensor): x
+        y (Number): The dtype is same as x.
+
+    Returns:
+        Tensor, has the same dtype as x.
+    """
+    return InplaceAddsExt()(x, y)
+
+
 @add.register("Tensor", "Number")
 def _tensor_add_scalar(x, y):
     """
@@ -165,6 +191,7 @@ def _tensor_add_scalar(x, y):
     return F.add(x, y)
 
 
+@augassign_add.register("Tuple", "Tensor")
 @add.register("Tuple", "Tensor")
 def _tuple_add_tensor(x, y):
     """
@@ -181,6 +208,7 @@ def _tuple_add_tensor(x, y):
     return F.tensor_add(x, y)
 
 
+@augassign_add.register("Tensor", "Tuple")
 @add.register("Tensor", "Tuple")
 def _tensor_add_tuple(x, y):
     """
@@ -197,6 +225,7 @@ def _tensor_add_tuple(x, y):
     return F.tensor_add(x, y)
 
 
+@augassign_add.register("List", "Tensor")
 @add.register("List", "Tensor")
 def _list_add_tensor(x, y):
     """
@@ -213,6 +242,7 @@ def _list_add_tensor(x, y):
     return F.tensor_add(x, y)
 
 
+@augassign_add.register("Tensor", "List")
 @add.register("Tensor", "List")
 def _tensor_add_list(x, y):
     """
@@ -229,6 +259,7 @@ def _tensor_add_list(x, y):
     return F.tensor_add(x, y)
 
 
+@augassign_add.register("List", "List")
 @add.register("List", "List")
 def _list_add_list(x, y):
     """
@@ -246,6 +277,21 @@ def _list_add_list(x, y):
     return _list_add(x, y)
 
 
+@augassign_add.register("Tensor", "Tensor")
+def _tensor_add_tensor_augassign(x, y):
+    """
+    Returns x + y element-wise.
+
+    Args:
+        x (Tensor): x
+        y (Tensor): The dtype is same as x.
+
+    Returns:
+        Tensor, has the same dtype as x.
+    """
+    return InplaceAddExt()(x, y)
+
+
 @add.register("Tensor", "Tensor")
 def _tensor_add_tensor(x, y):
     """
@@ -261,6 +307,7 @@ def _tensor_add_tensor(x, y):
     return F.add(x, y)
 
 
+@augassign_add.register("RowTensor", "Tensor")
 @add.register("RowTensor", "Tensor")
 def add_rowtensor_tensor(x, y):
     """
@@ -306,6 +353,7 @@ def _add_env(x, y):
     return F.environ_add(x, y)
 
 
+@augassign_add.register("Tuple", "Tuple")
 @add.register("Tuple", "Tuple")
 def _add_tuple(x, y):
     """
@@ -323,6 +371,7 @@ def _add_tuple(x, y):
     return _tuple_add(x, y)
 
 
+@augassign_add.register("CSRTensor", "CSRTensor")
 @add.register("CSRTensor", "CSRTensor")
 def _add_csrtensor(x, y):
     """
@@ -340,6 +389,7 @@ def _add_csrtensor(x, y):
     return F.csr_add(x, y, make_tensor(1, x.values.dtype), make_tensor(1, x.values.dtype))
 
 
+@augassign_add.register("COOTensor", "COOTensor")
 @add.register("COOTensor", "COOTensor")
 def _add_cootensor(x, y):
     """
@@ -357,6 +407,7 @@ def _add_cootensor(x, y):
     return F.coo_add(x, y, make_tensor(0, x.values.dtype))
 
 
+@augassign_add.register("COOTensor", "Tensor")
 @add.register("COOTensor", "Tensor")
 def _add_cootensor_tensor(x, y):
     """
@@ -374,6 +425,7 @@ def _add_cootensor_tensor(x, y):
     return F.tensor_scatter_add(y, x.indices, x.values)
 
 
+@augassign_add.register("Tensor", "COOTensor")
 @add.register("Tensor", "COOTensor")
 def _add_tensor_cootensor(x, y):
     """
@@ -451,6 +503,21 @@ def _add_rowtensor_tensor(x, y):
     return x + y
 
 
+@_add_backward.register("Tensor", "RowTensor")
+def _add_tensor_rowtensor(x, y):
+    """
+   Adds Tensor and RowTensor.
+
+   Args:
+       x (Tensor): x
+       y (RowTensor): y
+
+   Returns:
+       RowTensor, the dtype is same as y.
+   """
+    return y + x
+
+
 @_add_backward.register("None", "None")
 def _add_nonetensor_tensor(x, y):
     """
@@ -601,6 +668,7 @@ hyper_add = base.HyperMap(_add_backward)
 
 
 # pylint: disable=protected-access
+@augassign_add._register_default()
 @add._register_default()
 def default_add(x, y):
     """Default function for add."""

mindspore/ops/composite/multitype_ops/div_impl.py

@@ -16,6 +16,7 @@
 """Implementation for internal polymorphism `div` operations."""
 from __future__ import division
 
+from mindspore.ops.auto_generate.gen_ops_prim import InplaceDiv, InplaceDivs
 from mindspore.ops.composite.multitype_ops import _compile_utils as utils
 from mindspore.ops.composite.multitype_ops._constexpr_utils import log_warning, check_equal
 from mindspore.ops.composite import base
@@ -24,6 +25,15 @@ from mindspore.common import COOTensor
 from mindspore.ops.auto_generate import div_op
 
 
+# x /= y
+augassign_div = base.MultitypeFuncGraph("augassign_div", True)
+"""
+augassign_div is a metafuncgraph object which will div two objects according to input type
+using ".register" decorator
+"""
+augassign_div.set_need_raise()
+
+
 div = base.MultitypeFuncGraph("div", True)
 """
 div is a metafuncgraph object which will div two objects according to input type
@@ -32,6 +42,7 @@ using ".register" decorator
 div.set_need_raise()
 
 
+@augassign_div.register("CSRTensor", "Tensor")
 @div.register("CSRTensor", "Tensor")
 def _csrtensor_div_tensor(x, y):
     """
@@ -44,6 +55,7 @@ def _csrtensor_div_tensor(x, y):
     return F.csr_div(x, y)
 
 
+@augassign_div.register("COOTensor", "Tensor")
 @div.register("COOTensor", "Tensor")
 def _cootensor_div_tensor(x, y):
     """
@@ -58,6 +70,7 @@ def _cootensor_div_tensor(x, y):
     return COOTensor(x.indices, x.values / other_values, x.shape)
 
 
+@augassign_div.register("Number", "Number")
 @div.register("Number", "Number")
 def _div_scalar(x, y):
     """
@@ -73,6 +86,21 @@ def _div_scalar(x, y):
     return F.scalar_div(x, y)
 
 
+@augassign_div.register("Tensor", "Tensor")
+def _div_tensor_augassign(x, y):
+    """
+    Two tensors divide by element.
+
+    Args:
+        x (Tensor): The first input tensor.
+        y (Tensor): The second input tensor.
+
+    Returns:
+        Tensor, has the same dtype as x.
+    """
+    return InplaceDiv()(x, y)
+
+
 @div.register("Tensor", "Tensor")
 def _div_tensor(x, y):
     """
@@ -88,6 +116,7 @@ def _div_tensor(x, y):
     return div_op(x, y)
 
 
+@augassign_div.register("Number", "Tensor")
 @div.register("Number", "Tensor")
 def _scalar_div_tensor(x, y):
     """
@@ -103,6 +132,21 @@ def _scalar_div_tensor(x, y):
     return div_op(x, y)
 
 
+@augassign_div.register("Tensor", "Number")
+def _tensor_div_scalar_augassign(x, y):
+    """
+    Tensor divided by number.
+
+    Args:
+        x (Tensor): x
+        y (Number): The dtype is same as x.
+
+    Returns:
+        Tensor, has the same dtype as x.
+    """
+    return InplaceDivs()(x, y)
+
+
 @div.register("Tensor", "Number")
 def _tensor_div_scalar(x, y):
     """
@@ -118,6 +162,7 @@ def _tensor_div_scalar(x, y):
     return div_op(x, y)
 
 
+@augassign_div.register("Tuple", "Tensor")
 @div.register("Tuple", "Tensor")
 def _tuple_div_tensor(x, y):
     """
@@ -134,6 +179,7 @@ def _tuple_div_tensor(x, y):
     return F.tensor_div(x, y)
 
 
+@augassign_div.register("Tensor", "Tuple")
 @div.register("Tensor", "Tuple")
 def _tensor_div_tuple(x, y):
     """
@@ -150,6 +196,7 @@ def _tensor_div_tuple(x, y):
     return F.tensor_div(x, y)
 
 
+@augassign_div.register("List", "Tensor")
 @div.register("List", "Tensor")
 def _list_div_tensor(x, y):
     """
@@ -166,6 +213,7 @@ def _list_div_tensor(x, y):
     return F.tensor_div(x, y)
 
 
+@augassign_div.register("Tensor", "List")
 @div.register("Tensor", "List")
 def _tensor_div_list(x, y):
     """
@@ -183,6 +231,7 @@ def _tensor_div_list(x, y):
 
 
 # pylint: disable=protected-access
+@augassign_div._register_default()
 @div._register_default()
 def default_div(x, y):
     """Default function for div."""

mindspore/ops/composite/multitype_ops/floordiv_impl.py

@@ -15,11 +15,21 @@
 
 """Implementation for internal polymorphism `floordiv` operations."""
 from __future__ import absolute_import
+from mindspore.ops.auto_generate.gen_ops_prim import InplaceFloorDivide, InplaceFloorDivides
 from mindspore.ops.composite.multitype_ops import _compile_utils as utils
 from mindspore.ops.composite import base
 from mindspore.ops import functional as F
 
 
+# x //= y
+augassign_floordiv = base.MultitypeFuncGraph("augassign_floordiv", True)
+"""
+`augassign_floordiv` is a metafuncgraph object which will compute the floordiv of two objects
+using ".register" decorator.
+"""
+augassign_floordiv.set_need_raise()
+
+
 floordiv = base.MultitypeFuncGraph("floordiv", True)
 """
 `floordiv` is a metafuncgraph object which will compute the floordiv of two objects
@@ -28,30 +38,45 @@ using ".register" decorator.
 floordiv.set_need_raise()
 
 
+@augassign_floordiv.register("Number", "Number")
 @floordiv.register("Number", "Number")
 def _floordiv_scalar(x, y):
     """Returns x // y where x and y are all scalars."""
     return F.scalar_floordiv(x, y)
 
 
+@augassign_floordiv.register("Tensor", "Tensor")
+def _floordiv_tensor_augassign(x, y):
+    """Returns x // y where x and y are all tensors."""
+    return InplaceFloorDivide()(x, y)
+
+
 @floordiv.register("Tensor", "Tensor")
 def _floordiv_tensor(x, y):
     """Returns x // y where x and y are all tensors."""
     return F.tensor_floordiv(x, y)
 
 
+@augassign_floordiv.register("Tensor", "Number")
+def _tensor_floordiv_scalar_augassign(x, y):
+    """Returns x // y where x is a tensor and y is a scalar. x and y should have same dtype."""
+    return InplaceFloorDivides()(x, y)
+
+
 @floordiv.register("Tensor", "Number")
 def _tensor_floordiv_scalar(x, y):
     """Returns x // y where x is a tensor and y is a scalar. x and y should have same dtype."""
     return F.tensor_floordiv(x, y)
 
 
+@augassign_floordiv.register("Number", "Tensor")
 @floordiv.register("Number", "Tensor")
 def _scalar_floordiv_tensor(x, y):
     """Returns x // y where x is a scalar and y is a tensor. x and y should have same dtype."""
     return F.tensor_floordiv(x, y)
 
 
+@augassign_floordiv.register("Tuple", "Tensor")
 @floordiv.register("Tuple", "Tensor")
 def _tuple_floordiv_tensor(x, y):
     """Returns x // y where x is a tuple and y is a tensor. """
@@ -59,6 +84,7 @@ def _tuple_floordiv_tensor(x, y):
     return F.tensor_floordiv(x, y)
 
 
+@augassign_floordiv.register("Tensor", "Tuple")
 @floordiv.register("Tensor", "Tuple")
 def _tensor_floordiv_tuple(x, y):
     """Returns x // y where x is a tensor and y is a tuple. """
@@ -66,6 +92,7 @@ def _tensor_floordiv_tuple(x, y):
     return F.tensor_floordiv(x, y)
 
 
+@augassign_floordiv.register("List", "Tensor")
 @floordiv.register("List", "Tensor")
 def _list_floordiv_tensor(x, y):
     """Returns x // y where x is a list and y is a tensor. """
@@ -73,6 +100,7 @@ def _list_floordiv_tensor(x, y):
     return F.tensor_floordiv(x, y)
 
 
+@augassign_floordiv.register("Tensor", "List")
 @floordiv.register("Tensor", "List")
 def _tensor_floordiv_list(x, y):
     """Returns x // y where x is a tensor and y is a list. """
@@ -81,6 +109,7 @@ def _tensor_floordiv_list(x, y):
 
 
 # pylint: disable=protected-access
+@augassign_floordiv._register_default()
 @floordiv._register_default()
 def default_floordiv(x, y):
     """Default function for floordiv."""

mindspore/ops/composite/multitype_ops/getitem_impl.py

@@ -15,6 +15,7 @@
 
 """Implementation for internal polymorphism `getitem` operations."""
 from __future__ import absolute_import
+import os
 from mindspore.ops.operations import _inner_ops as inner
 from mindspore.ops.operations import _map_tensor_ops
 from mindspore.ops.composite.multitype_ops import _compile_utils as compile_utils
@@ -286,6 +287,8 @@ def _tensor_getitem_by_number(data, number_index):
     Outputs:
         Tensor, element type is as same as the element type of data.
     """
+    if os.environ.get("MS_DEV_TENSOR_INDEX_BOOST") == "1":
+        return compile_utils._tensor_getitem(data, number_index)
     return compile_utils.tensor_index_by_number(data, number_index)
 
 
@@ -316,6 +319,8 @@ def _tensor_getitem_by_slice(data, slice_index):
     Outputs:
         Tensor, element type is the same as the element type of data.
     """
+    if os.environ.get("MS_DEV_TENSOR_INDEX_BOOST") == "1":
+        return compile_utils._tensor_getitem(data, slice_index)
     return compile_utils.tensor_index_by_slice(data, slice_index)
 
 
@@ -331,6 +336,8 @@ def _tensor_getitem_by_tensor(data, tensor_index):
     Outputs:
         Tensor, element type is the same as the element type of data.
     """
+    if os.environ.get("MS_DEV_TENSOR_INDEX_BOOST") == "1":
+        return compile_utils._tensor_getitem(data, tensor_index)
     return compile_utils.tensor_index_by_tensor(data, tensor_index)
 
 
@@ -361,6 +368,8 @@ def _tensor_getitem_by_list(data, list_index):
     Outputs:
         Tensor ,same as data.
     """
+    if os.environ.get("MS_DEV_TENSOR_INDEX_BOOST") == "1":
+        return compile_utils._tensor_getitem(data, list_index)
     return compile_utils.tensor_index_by_list(data, list_index)
 
 
@@ -376,6 +385,8 @@ def _tensor_getitem_by_tuple(data, tuple_index):
     Outputs:
         Tensor, element type is the same as the element type of data.
     """
+    if os.environ.get("MS_DEV_TENSOR_INDEX_BOOST") == "1":
+        return compile_utils._tensor_getitem(data, tuple_index)
     return compile_utils.tensor_index_by_tuple(data, tuple_index)
 
 

mindspore/ops/composite/multitype_ops/mod_impl.py

@@ -18,6 +18,8 @@ from __future__ import absolute_import
 from mindspore.ops.composite.multitype_ops import _compile_utils as utils
 from mindspore.ops.composite import base
 from mindspore.ops import functional as F
+from mindspore.ops.auto_generate import (remainder_tensor_tensor_op, remainder_tensor_scalar_op,
+                                         remainder_scalar_tensor_op)
 
 
 mod = base.MultitypeFuncGraph("mod", True)
@@ -37,19 +39,19 @@ def _mod_scalar(x, y):
 @mod.register("Tensor", "Tensor")
 def _mod_tensor(x, y):
     """Returns x % y where x and y are all tensors."""
-    return F.tensor_mod(x, y)
+    return remainder_tensor_tensor_op(x, y)
 
 
 @mod.register("Tensor", "Number")
 def _tensor_mod_scalar(x, y):
     """Returns x % y where x is a tensor and y is a scalar. x and y should have same dtype."""
-    return F.tensor_mod(x, y)
+    return remainder_tensor_scalar_op(x, y)
 
 
 @mod.register("Number", "Tensor")
 def _scalar_mod_tensor(x, y):
     """Returns x % y where x is a scalar and y is a tensor. x and y should have same dtype."""
-    return F.tensor_mod(x, y)
+    return remainder_scalar_tensor_op(x, y)
 
 
 @mod.register("Tuple", "Tensor")

mindspore/ops/composite/multitype_ops/mul_impl.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """Implementation for internal polymorphism `mul` operations."""
 from mindspore.ops.operations import _inner_ops as inner
+from mindspore.ops.auto_generate.gen_ops_prim import InplaceMul, InplaceMuls
 from mindspore.ops.composite.multitype_ops import _compile_utils as utils
 from mindspore.ops.composite.multitype_ops._constexpr_utils import check_equal
 from mindspore.ops.composite import base
@@ -22,6 +23,15 @@ from mindspore.common import COOTensor
 from ...operations._sequence_ops import SequenceMul
 
 
+# x *= y
+augassign_mul = base.MultitypeFuncGraph("augassign_mul", True)
+"""
+`augassign_mul` is a metafuncgraph object which will multiply two objects according to input type
+using ".register" decorator.
+"""
+augassign_mul.set_need_raise()
+
+
 mul = base.MultitypeFuncGraph("mul", True)
 """
 `mul` is a metafuncgraph object which will multiply two objects according to input type
@@ -30,6 +40,7 @@ using ".register" decorator.
 mul.set_need_raise()
 
 
+@augassign_mul.register("Number", "Number")
 @mul.register("Number", "Number")
 def _mul_scalar(x, y):
     """
@@ -41,6 +52,17 @@ def _mul_scalar(x, y):
     return F.scalar_mul(x, y)
 
 
+@augassign_mul.register("Tensor", "Tensor")
+def _mul_tensor_augassign(x, y):
+    """
+    Returns x * y by element-wise where x and y are all tensors.
+
+    Outputs:
+        Tensor, has the same dtype as x.
+    """
+    return InplaceMul()(x, y)
+
+
 @mul.register("Tensor", "Tensor")
 def _mul_tensor(x, y):
     """
@@ -52,6 +74,7 @@ def _mul_tensor(x, y):
     return F.tensor_mul(x, y)
 
 
+@augassign_mul.register("Number", "Tensor")
 @mul.register("Number", "Tensor")
 def _scalar_mul_tensor(x, y):
     """
@@ -63,6 +86,17 @@ def _scalar_mul_tensor(x, y):
     return F.tensor_mul(x, y)
 
 
+@augassign_mul.register("Tensor", "Number")
+def _tensor_mul_scalar_augassign(x, y):
+    """
+    Returns x * y where x is a tensor and y is a scalar. x and y have same dtype.
+
+    Outputs:
+        Tensor, has the same dtype as x.
+    """
+    return InplaceMuls()(x, y)
+
+
 @mul.register("Tensor", "Number")
 def _tensor_mul_scalar(x, y):
     """
@@ -74,6 +108,7 @@ def _tensor_mul_scalar(x, y):
     return F.tensor_muls(x, y)
 
 
+@augassign_mul.register("Number", "String")
 @mul.register("Number", "String")
 def _number_mul_string(x, y):
     """
@@ -85,6 +120,7 @@ def _number_mul_string(x, y):
     return inner.string_mul(x, y)
 
 
+@augassign_mul.register("String", "Number")
 @mul.register("String", "Number")
 def _string_mul_number(x, y):
     """
@@ -96,6 +132,7 @@ def _string_mul_number(x, y):
     return inner.string_mul(x, y)
 
 
+@augassign_mul.register("List", "Number")
 @mul.register("List", "Number")
 def _list_mul_scalar(x, y):
     """
@@ -116,6 +153,7 @@ def _list_mul_scalar(x, y):
     return res
 
 
+@augassign_mul.register("Number", "List")
 @mul.register("Number", "List")
 def _scalar_mul_list(x, y):
     """
@@ -136,6 +174,7 @@ def _scalar_mul_list(x, y):
     return res
 
 
+@augassign_mul.register("Tuple", "Number")
 @mul.register("Tuple", "Number")
 def _tuple_mul_scalar(x, y):
     """
@@ -156,6 +195,7 @@ def _tuple_mul_scalar(x, y):
     return res
 
 
+@augassign_mul.register("Number", "Tuple")
 @mul.register("Number", "Tuple")
 def _scalar_mul_tuple(x, y):
     """
@@ -176,6 +216,7 @@ def _scalar_mul_tuple(x, y):
     return res
 
 
+@augassign_mul.register("Tensor", "Tuple")
 @mul.register("Tensor", "Tuple")
 def _tensor_mul_tuple(x, y):
     """
@@ -192,6 +233,7 @@ def _tensor_mul_tuple(x, y):
     return F.tensor_mul(x, y)
 
 
+@augassign_mul.register("Tuple", "Tensor")
 @mul.register("Tuple", "Tensor")
 def _tuple_mul_tensor(x, y):
     """
@@ -208,6 +250,7 @@ def _tuple_mul_tensor(x, y):
     return F.tensor_mul(x, y)
 
 
+@augassign_mul.register("Tensor", "List")
 @mul.register("Tensor", "List")
 def _tensor_mul_list(x, y):
     """
@@ -224,6 +267,7 @@ def _tensor_mul_list(x, y):
     return F.tensor_mul(x, y)
 
 
+@augassign_mul.register("List", "Tensor")
 @mul.register("List", "Tensor")
 def _list_mul_tensor(x, y):
     """
@@ -240,6 +284,7 @@ def _list_mul_tensor(x, y):
     return F.tensor_mul(x, y)
 
 
+@augassign_mul.register("CSRTensor", "Tensor")
 @mul.register("CSRTensor", "Tensor")
 def _csrtensor_mul_tensor(x, y):
     """
@@ -251,6 +296,7 @@ def _csrtensor_mul_tensor(x, y):
     return F.csr_mul(x, y)
 
 
+@augassign_mul.register("Tensor", "CSRTensor")
 @mul.register("Tensor", "CSRTensor")
 def _tensor_mul_csrtensor(x, y):
     """
@@ -262,6 +308,7 @@ def _tensor_mul_csrtensor(x, y):
     return F.csr_mul(y, x)
 
 
+@augassign_mul.register("COOTensor", "Tensor")
 @mul.register("COOTensor", "Tensor")
 def _cootensor_mul_tensor(x, y):
     """
@@ -275,6 +322,7 @@ def _cootensor_mul_tensor(x, y):
     return COOTensor(x.indices, x.values * other_values, x.shape)
 
 
+@augassign_mul.register("Tensor", "COOTensor")
 @mul.register("Tensor", "COOTensor")
 def _tensor_mul_cootensor(x, y):
     """
@@ -289,6 +337,7 @@ def _tensor_mul_cootensor(x, y):
 
 
 # pylint: disable=protected-access
+@augassign_mul._register_default()
 @mul._register_default()
 def default_mul(x, y):
     """Default function for mul."""