dev(hansbug): use stream function to wrap native matmul function

treetensor/torch/__init__.py

@@ -10,6 +10,8 @@ from .funcs import __all__ as _funcs_all
 from .funcs.base import get_func_from_torch
 from .size import *
 from .size import __all__ as _size_all
+from .stream import *
+from .stream import __all__ as _stream_all
 from .tensor import *
 from .tensor import __all__ as _tensor_all
 from ..config.meta import __VERSION__
@@ -18,6 +20,7 @@ __all__ = [
     *_funcs_all,
     *_size_all,
     *_tensor_all,
+    *_stream_all,
 ]
 
 _basic_types = (

treetensor/torch/funcs/matrix.py

 import torch
 
 from .base import doc_from_base, func_treelize
+from ..stream import stream_call
 
 __all__ = [
     'dot', 'matmul', 'mm',
@@ -73,7 +74,7 @@ def matmul(input, other, *args, **kwargs):
         └── b --> <Tensor 0x7f2e74886430>
             └── x --> tensor(40)
     """
-    return torch.matmul(input, other, *args, **kwargs)
+    return stream_call(torch.matmul, input, other, *args, **kwargs)
 
 
 # noinspection PyShadowingBuiltins