fix jittor.nn.AdaptiveMaxPool3d doc

README.md

@@ -382,10 +382,10 @@ Email: [email protected]
 
 File an issue: https://github.com/Jittor/jittor/issues
 
-QQ Group: 761222083
+QQ Group: 836860279
 
 
-<img src="https://cg.cs.tsinghua.edu.cn/jittor/images/news/2020-12-8-21-19-1_2_2/fig4.png" width="200"/>
+<img src="https://github.com/Jittor/jittor/assets/62846124/8dd830bd-b31c-4e4f-9a78-5fd7a3409145" width="200"/>
 
 ## The Team
 

doc/source/jittor.nn.md

@@ -10,7 +10,7 @@ jittor.nn
 
 .. automodule:: jittor.nn
    :imported-members:
-   :members: Pool, pool, AdaptiveAvgPool2d, Pool3d, AdaptiveMaxPool2d, AdaptiveAvgPool3d, AdaptiveMaxPool2d, pool3d, AvgPool2d, AvgPool3d, avg_pool2d, MaxPool2d, MaxPool3d, max_pool2d, max_pool3d, MaxUnpool2d, MaxUnpool3d
+   :members: Pool, pool, AdaptiveAvgPool2d, Pool3d, AdaptiveMaxPool2d, AdaptiveAvgPool3d, AdaptiveMaxPool3d, pool3d, AvgPool2d, AvgPool3d, avg_pool2d, MaxPool2d, MaxPool3d, max_pool2d, max_pool3d, MaxUnpool2d, MaxUnpool3d
    :undoc-members: 
 
 .. autoclass:: jittor.nn.ReLU

python/jittor/__init__.py

@@ -9,7 +9,7 @@
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
 
-__version__ = '1.3.9.6'
+__version__ = '1.3.9.10'
 from jittor_utils import lock
 with lock.lock_scope():
     ori_int = int
@@ -428,7 +428,9 @@ def random(shape, dtype="float32", type="uniform"):
         jt.Var([[0.96788853 0.28334728 0.30482838]
                 [0.46107793 0.62798643 0.03457401]], dtype=float32)
     '''
-
+    for dim in shape:
+        if dim < 0:
+            raise RuntimeError(f"Trying to create tensor with negative dimension {dim}: {shape}")
     ret = ops.random(shape, "float32", type)
    ## TODO: move those code to core
    #if dtype in ["float16", "bfloat16"]:
@@ -484,6 +486,9 @@ def ones(*shape, dtype="float32"):
         shape = shape[:-1]
     if isinstance(shape, tuple) and isinstance(shape[0], (Sequence, NanoVector)):
         shape = shape[0]
+    for dim in shape:
+        if dim < 0:
+            raise RuntimeError(f"Trying to create tensor with negative dimension {dim}: {shape}")
     return unary(1, dtype).broadcast(shape)
 
 def new_ones(x, size):
@@ -515,6 +520,9 @@ def zeros(*shape, dtype="float32"):
         shape = shape[:-1]
     if isinstance(shape, tuple) and isinstance(shape[0], (Sequence, NanoVector)):
         shape = shape[0]
+    for dim in shape:
+        if dim < 0:
+            raise RuntimeError(f"Trying to create tensor with negative dimension {dim}: {shape}")
     return unary(0, dtype).broadcast(shape)
 
 def new_zeros(x, size):
@@ -547,6 +555,9 @@ def full(shape,val,dtype="float32"):
     '''
     if not isinstance(shape, (NanoVector, Sequence)):
         shape = (shape,)
+    for dim in shape:
+        if dim < 0:
+            raise RuntimeError(f"Trying to create tensor with negative dimension {dim}: {shape}")
     return unary(val, dtype).broadcast(shape)
 
 def new_full(x, size, val):
@@ -641,14 +652,22 @@ def var(x, dim=None, dims=None, unbiased=False, keepdims=False):
     return sqr
 Var.var = var
 
-def std(x):
-    matsize=1
-    for i in x.shape:
-        matsize *= i
-    out=(x-x.mean()).sqr().sum()
-    out=out/(matsize-1)
-    out=out.maximum(1e-6).sqrt()
-    return out
+def std(x, dim=None, keepdim=False):
+    if dim is None:
+        matsize=1
+        for i in x.shape:
+            matsize *= i
+        out=(x-x.mean()).sqr().sum()
+        out=out/(matsize-1)
+        out=out.maximum(1e-6).sqrt()
+        return out
+    else:
+        dimsize=x.size(dim)
+        mean=jt.mean(x, dim, keepdim=True)
+        out=(x - mean).sqr().sum(dim=dim, keepdim=keepdim)
+        out=out/(dimsize-1)
+        out=out.maximum(1e-6).sqrt()
+        return out
 Var.std = std
 
 def norm(x, p=2, dim=-1, keepdims=False, eps=1e-30, keepdim=False):
@@ -687,6 +706,8 @@ def flatten(input, start_dim=0, end_dim=-1):
     start_dim = len(in_shape) + start_dim if start_dim < 0 else start_dim
     end_dim = len(in_shape) + end_dim if end_dim < 0 else end_dim
     assert end_dim >= start_dim, "end_dim should be larger than or equal to start_dim for flatten function"
+    if len(in_shape) <= end_dim:
+        raise IndexError(f"Dimension out of range (expected to be in range of [{-len(in_shape)}, {len(in_shape) - 1}], but got {end_dim})")
     out_shape = []
     for i in range(0,start_dim,1): out_shape.append(in_shape[i])
     dims = 1
@@ -917,6 +938,9 @@ def randn(*size, dtype="float32", requires_grad=True) -> Var:
          [-0.612632   -1.1471151  -1.1879086 ]], dtype=float32)
     '''
     if isinstance(size, tuple) and isinstance(size[0], (tuple, list, NanoVector)): size = size[0]
+    for dim in size:
+        if dim < 0:
+            raise RuntimeError(f"Trying to create tensor with negative dimension {dim}: {size}")
     arr = jt.random(size, dtype, "normal")
     if not requires_grad: return arr.stop_grad()
     return arr
@@ -1013,6 +1037,9 @@ def randint(low, high=None, shape=(1,), dtype="int32") -> Var:
                 [1 1 1]], dtype=int32)
     '''
     if high is None: low, high = 0, low
+    for dim in shape:
+        if dim < 0:
+            raise RuntimeError(f"Trying to create tensor with negative dimension {dim}: {shape}")
     v = (jt.random(shape) * (high - low) + low).clamp(low, high-0.5)
     v = jt.floor_int(v)
     return v.astype(dtype)
@@ -1437,9 +1464,17 @@ def requires_grad_(self, requires_grad=True):
     def __hooked_call__(self, *args, **kw):
         if hasattr(self, "__fhook2__"):
             if len(kw):
-                self.__fhook2__(self, args, kw)
+                args_kw_result = self.__fhook2__(self, args, kw)
             else:
-                self.__fhook2__(self, args)
+                args_kw_result = self.__fhook2__(self, args)
+            if args_kw_result is not None:
+                if isinstance(args_kw_result, tuple) and len(args_kw_result) == 2:
+                    args, kw = args_kw_result
+                else:
+                    raise RuntimeError(
+                        "forward pre-hook must return None or a tuple "
+                        f"of (new_args, new_kwargs), but got {args_kw_result}."
+                    )
         if hasattr(self, "__bihook__"):
             if len(kw):
                 LOG.w("backward hook not support kw")
@@ -1458,9 +1493,11 @@ def __hooked_call__(self, *args, **kw):
                 ret = grad_hooker(ret, self.__bohook__)
         if hasattr(self, "__fhook__"):
             if len(kw):
-                self.__fhook__(self, args, ret, kw)
+                res = self.__fhook__(self, args, ret, kw)
             else:
-                self.__fhook__(self, args, ret)
+                res = self.__fhook__(self, args, ret)
+            if res is not None:
+                ret = res
         return ret
 
     def _place_hooker(self):
@@ -1595,6 +1632,8 @@ def load_parameters(self, params):
                 else:
                     if hasattr(v, k):
                         v = getattr(v, k)
+                        if v is None:
+                            continue
                         assert isinstance(v, (Module, Var)), \
                             f"expect a jittor Module or Var, but got <{v.__class__.__name__}>, key: {key}"
                     else:
@@ -2119,6 +2158,7 @@ def is_var(v):
 from . import optim
 from . import dataset
 from . import init
+from . import gradfunctional
 
 dtype = NanoString
 
@@ -2152,3 +2192,7 @@ def inplace_wrapper(new_k, prev_func):
 from . import math_util
 from .math_util import *
 from . import distributions
+
+if jt.compiler.has_acl:
+    from jittor.extern.acl.acl_compiler import change_function
+    change_function()

python/jittor/__init__.pyi

@@ -1,7 +1,7 @@
 from jittor_core import *
 from jittor_core.ops import *
 from .misc import *
-from . import attention as attention, contrib as contrib, dataset as dataset, init as init, linalg as linalg, lr_scheduler as lr_scheduler, numpy2cupy as numpy2cupy, optim as optim, sparse as sparse
+from . import attention as attention, contrib as contrib, dataset as dataset, init as init, linalg as linalg, lr_scheduler as lr_scheduler, numpy2cupy as numpy2cupy, optim as optim, sparse as sparse, gradfunctional as gradfunctional
 from .compile_extern import cublas as cublas, cudnn as cudnn, cufft as cufft, curand as curand, mkl_ops as mkl_ops, mpi_ops as mpi_ops, world_size as world_size
 from .compiler import compile_custom_op as compile_custom_op, compile_custom_ops as compile_custom_ops
 from .contrib import concat as concat