version 1.2

Author: cjld

extern/cuda/cub/ops/cub_arg_reduce_op.cc

@@ -49,6 +49,8 @@ void CubArgReduceOp::infer_shape() {
     if (keepdims) {
         shape.push_back(1);
     }
+    if (shape.size() == 0)
+        shape.push_back(1);
     y->set_shape(shape);
     y_key->set_shape(shape);
 }
@@ -104,4 +106,4 @@ void CubArgReduceOp::jit_run() {
 #endif // JIT_cuda
 #endif // JIT
 
-} // jittor
+} // jittor

extern/cuda/cudnn/ops/cudnn_conv_backward_w_op.cc

@@ -42,8 +42,8 @@ static inline void set_shape(Var* x, const char* f, const string& format, int a,
         shape[0], shape[1], shape[2], shape[3]));
 }
 
-CudnnConvBackwardWOp::CudnnConvBackwardWOp(Var* x, Var* dy, int kernel_size, int stride, int padding, int dilation, int groups, string xformat, string wformat, string yformat)
-        : x(x), dy(dy), kernel_size(kernel_size), stride(stride), padding(padding), dilation(dilation), groups(groups),
+CudnnConvBackwardWOp::CudnnConvBackwardWOp(Var* x, Var* dy, int kh, int kw, int stride, int padding, int dilation, int groups, string xformat, string wformat, string yformat)
+        : x(x), dy(dy), kh(kh), kw(kw), stride(stride), padding(padding), dilation(dilation), groups(groups),
       xformat(move(xformat)), wformat(move(wformat)), yformat(move(yformat)) {
     flags.set(NodeFlags::_cuda, 1);
     flags.set(NodeFlags::_cpu, 0);
@@ -57,8 +57,8 @@ void CudnnConvBackwardWOp::infer_shape() {
     get_shape(x, "abcd", xformat, xn, xc, xh, xw);
     get_shape(dy, "abcd", yformat, yn, yc, yh, yw);
     wco = yc, wci = xc / groups;
-    wh = kernel_size;
-    ww = kernel_size;
+    wh = kh;
+    ww = kw;
     set_shape(dw, "oihw", wformat, wco, wci, wh, ww);
 }
 

extern/cuda/cudnn/ops/cudnn_conv_backward_w_op.h

@@ -13,14 +13,14 @@ namespace jittor {
 
 struct CudnnConvBackwardWOp : Op {
     Var* x, * dy, * dw;
-    int kernel_size, stride, padding, dilation, groups;
+    int kh, kw, stride, padding, dilation, groups;
     string xformat, wformat, yformat;
 
-    CudnnConvBackwardWOp(Var* x, Var* y, int kernel_size, int stride, int padding, int dilation, int groups=1, string xformat="abcd", string wformat="oihw", string yformat="abcd");
+    CudnnConvBackwardWOp(Var* x, Var* y, int kh, int kw, int stride, int padding, int dilation, int groups=1, string xformat="abcd", string wformat="oihw", string yformat="abcd");
 
     const char* name() const override { return "cudnn_conv_backward_w"; }
     void infer_shape() override;
     DECLARE_jit_run;
 };
 
-} // jittor
+} // jittor

extern/mkl/ops/mkl_conv_backward_w_op.cc

@@ -45,8 +45,8 @@ static inline void set_shape(Var* x, const char* f, const string& format, int a,
         shape[0], shape[1], shape[2], shape[3]));
 }
 
-MklConvBackwardWOp::MklConvBackwardWOp(Var* x, Var* dy, int kernel_size, int stride, int padding, int dilation, int groups, string xformat, string wformat, string yformat)
-        : x(x), dy(dy), kernel_size(kernel_size), stride(stride), padding(padding), dilation(dilation), groups(groups), 
+MklConvBackwardWOp::MklConvBackwardWOp(Var* x, Var* dy, int kh, int kw, int stride, int padding, int dilation, int groups, string xformat, string wformat, string yformat)
+        : x(x), dy(dy), kh(kh), kw(kw), stride(stride), padding(padding), dilation(dilation), groups(groups), 
       xformat(move(xformat)), wformat(move(wformat)), yformat(move(yformat)) {
     dw = create_output(nullptr, dtype_infer(dy->ns, x->ns));
 }
@@ -58,8 +58,8 @@ void MklConvBackwardWOp::infer_shape() {
     get_shape(x, "abcd", xformat, xn, xc, xh, xw);
     get_shape(dy, "abcd", yformat, yn, yc, yh, yw);
     wco = yc, wci = xc / groups;
-    wh = kernel_size;
-    ww = kernel_size;
+    wh = kh;
+    ww = kw;
     set_shape(dw, "oihw", wformat, wco, wci, wh, ww);
 }
 
@@ -97,7 +97,8 @@ void MklConvBackwardWOp::jit_run() {
     int height = x->shape[findc("@XFORMAT",'c')];
     int width = x->shape[findc("@XFORMAT",'d')];
     int ch_out = dw->shape[findc("@WFORMAT",'o')];
-    int kernel_size = dw->shape[findc("@WFORMAT",'h')];
+    int kh = dw->shape[findc("@WFORMAT",'h')];
+    int kw = dw->shape[findc("@WFORMAT",'w')];
 
     auto* __restrict__ net_src = x->ptr<Txd>();
     auto* __restrict__ net_diff_dst = dy->ptr<Tyd>();
@@ -114,9 +115,9 @@ void MklConvBackwardWOp::jit_run() {
 
     memory::dims conv_src_tz = {batch, ch_in, height, width};
     memory::dims conv_weights_tz = groups>1
-        ? memory::dims{groups, ch_out/groups, ch_in/groups, kernel_size, kernel_size} 
-        : memory::dims{ch_out, ch_in, kernel_size, kernel_size};
-    memory::dims conv_dst_tz = {batch, ch_out, (height+padding*2-kernel_size*dilation+dilation-1)/stride+1, (width+padding*2-kernel_size*dilation+dilation-1)/stride+1};
+        ? memory::dims{groups, ch_out/groups, ch_in/groups, kh, kw} 
+        : memory::dims{ch_out, ch_in, kh, kw};
+    memory::dims conv_dst_tz = {batch, ch_out, (height+padding*2-kh*dilation+dilation-1)/stride+1, (width+padding*2-kw*dilation+dilation-1)/stride+1};
     memory::dims conv_strides = {stride, stride};
     memory::dims conv_padding = {padding, padding};
     memory::dims conv_dilation = {dilation-1, dilation-1};

extern/mkl/ops/mkl_conv_backward_w_op.h

@@ -13,14 +13,14 @@ namespace jittor {
 
 struct MklConvBackwardWOp : Op {
     Var* x, * dy, * dw;
-    int kernel_size, stride, padding, dilation, groups;
+    int kh, kw, stride, padding, dilation, groups;
     string xformat, wformat, yformat;
 
-    MklConvBackwardWOp(Var* x, Var* y, int kernel_size, int stride, int padding, int dilation, int groups=1, string xformat="abcd", string wformat="oihw", string yformat="abcd");
+    MklConvBackwardWOp(Var* x, Var* y, int kh, int kw, int stride, int padding, int dilation, int groups=1, string xformat="abcd", string wformat="oihw", string yformat="abcd");
 
     const char* name() const override { return "mkl_conv_backward_w"; }
     void infer_shape() override;
     DECLARE_jit_run;
 };
 
-} // jittor
+} // jittor

python/jittor/__init__.py

@@ -7,7 +7,7 @@
 # This file is subject to the terms and conditions defined in
 # file 'LICENSE.txt', which is part of this source code package.
 # ***************************************************************
-__version__ = '1.1.7.20'
+__version__ = '1.2.0.0'
 from . import lock
 with lock.lock_scope():
     from . import compiler
@@ -233,11 +233,22 @@ def ones(shape, dtype="float32"):
         shape = (shape,)
     return unary(1, dtype).broadcast(shape)
 
+def ones_like(x):
+    return ones(x.shape,x.dtype)
+
 def zeros(shape, dtype="float32"):
     if not isinstance(shape, (NanoVector, Sequence)):
         shape = (shape,)
     return unary(0, dtype).broadcast(shape)
 
+def full(shape,val,dtype="float32"):
+    if not isinstance(shape, (NanoVector, Sequence)):
+        shape = (shape,)
+    return unary(val, dtype).broadcast(shape)
+
+def zeros_like(x):
+    return zeros(x.shape,x.dtype)
+
 flags = core.flags()
 
 def std(x):
@@ -311,9 +322,17 @@ def squeeze(x, dim):
     return x.reshape(shape[:dim] + shape[dim+1:])
 Var.squeeze = squeeze
 
-def clamp(x, min_v, max_v):
-    assert min_v <= max_v
-    return x.maximum(min_v).minimum(max_v)
+def clamp(x, min_v=None, max_v=None):
+    if x.shape[0]==0:
+        return x
+    if min_v is not None and max_v is not None:
+        assert min_v <= max_v
+    if min_v is not None:
+        x = x.maximum(min_v)
+    if max_v is not None:
+        x = x.minimum(max_v)
+    return x
+
 Var.clamp = clamp
 
 def type_as(a, b):
@@ -574,6 +593,8 @@ def load_parameters(self, params):
                 else:
                     if hasattr(v, k):
                         v = getattr(v, k)
+                        assert isinstance(v, (Module, Var)), \
+                            f"expect a jittor Module or Var, but got <{v.__class__.__name__}>, key: {key}"
                     else:
                         end = 1
                         break
@@ -582,6 +603,8 @@ def load_parameters(self, params):
                     n_failed += 1
                     LOG.w(f'load parameter {key} failed ...')
             else:
+                assert isinstance(v, Var), \
+                    f"expect a jittor Var, but got <{v.__class__.__name__}>, key: {key}"
                 LOG.v(f'load parameter {key} success ...')
                 if isinstance(params[key], np.ndarray) or isinstance(params[key], list):
                     v.update(array(params[key]))
@@ -872,4 +895,4 @@ def to_bool(v):
 from . import contrib
 from . import numpy2cupy
 from .contrib import concat
-from .misc import *
+from .misc import *

python/jittor/compiler.py

@@ -241,7 +241,7 @@ def add_src(
         if "multiple_outputs" not in attrs:
             jit_cc_src.append(f"""
             VarPtr make_{cc_func_name}({", ".join(cc_make_args)}) {{
-                Op* _op = new {op_name}({", ".join(op_make_args)});
+                auto _op = new {op_name}({", ".join(op_make_args)});
                 if (_op->outputs_holder.size() != 1) {{
                     delete _op;
                     LOGf << "Wrong output size of" << \"{op_name}\";
@@ -261,7 +261,7 @@ def add_src(
         else:
             jit_cc_src.append(f"""
             vector<VarPtr> make_{cc_func_name}({", ".join(cc_make_args)}) {{
-                Op* _op = new {op_name}({", ".join(op_make_args)});
+                auto _op = new {op_name}({", ".join(op_make_args)});
                 if (_op->flags.get(NodeFlags::_forwarded)) {{
                     vector<VarPtr> outputs = move(_op->outputs_holder);
                     delete _op;
@@ -408,6 +408,15 @@ def add_src(
                         arg_type.replace("Var", "VarHolder")+' '+arg)
                     new_args.append(arg)
                     more_src.append(f"_op->add_inputs({arg});")
+                elif arg_type.startswith("VarSlices"):
+                    new_args_def.append(arg_def)
+                    new_args.append(arg)
+                    more_src.append(f"""
+                        vector<Var*> svars;
+                        for (int i=0; i<_op->vs.n; i++)
+                            if (_op->vs.slices[i].is_var())
+                                svars.push_back(_op->vs.slices[i].var);
+                        _op->add_inputs(svars);""")
                 else:
                     new_args_def.append(arg_def)
                     new_args.append(arg)

python/jittor/contrib.py

@@ -42,7 +42,7 @@ def concat(arr, dim):
         indexes[dim] = f"i{dim}-{cdim}"
         b = a.reindex(shape, indexes)
         # ugly fix for preventing large fused op 
-        if len(arr)>=10:
+        if len(arr)>=100:
             b.stop_fuse()
         if s is None:
             s = b
@@ -99,6 +99,20 @@ def slice_var_index(x, slices):
     cnt_list = 0
     extras_idx = []
     extras = []
+    has_ellipse = 0
+    ellipse_index = 0
+    for s,i in zip(slices,range(len(slices))):
+        if isinstance(s,type(...)):
+            has_ellipse+=1
+            ellipse_index = i
+    if has_ellipse>1:
+        raise Exception(f"There are more than one ...")
+    elif has_ellipse==1:
+        slices = list(slices)
+        del slices[ellipse_index]
+        while len(slices)<len(shape):
+            slices.insert(ellipse_index,slice(None))
+
     for i in range(len(shape)):
         if i>=len(slices):
             s = slice(None)
@@ -119,6 +133,7 @@ def slice_var_index(x, slices):
             step = 1 if s.step is None else s.step
             if start<0: start += sp
             if stop<0: stop += sp
+            if stop>sp+1: stop = sp
             out_shape.append(1+int(max(0, (stop-start-1)//step)))
             out_index.append(f"{start}+i{j}*{step}")
         elif isinstance(s, jt.Var):
@@ -160,3 +175,57 @@ def setitem(x, slices, value):
 
 jt.Var.__getitem__ = jt.Var.slice_var = slice_var
 jt.Var.__setitem__ = setitem
+
+# PATCH
+def getitem(x, slices):
+    if isinstance(slices, jt.Var) and slices.dtype == "bool":
+        return getitem(x, slices.where())
+    if isinstance(slices, list):
+        slices = tuple(slices)
+    return x.getitem(slices)
+
+def setitem(x, slices, value):
+    if isinstance(slices, jt.Var) and slices.dtype == "bool":
+        mask = jt.broadcast(slices, x)
+        value = jt.broadcast(value, x)
+        return mask.ternary(value, mask)
+    if isinstance(slices, list):
+        slices = tuple(slices)
+    return x.assign(x.setitem(slices, value))
+
+jt.Var.__getitem__ = jt.Var.slice_var = getitem
+jt.Var.__setitem__ = setitem
+
+def concat(arr, dim):
+    '''Concat Operator can concat a list of jt Var at a specfic dimension.
+    
+    * [in] x:   input var list for concat
+
+    * [in] dim: concat which dim
+
+    * [out] out:  concat result
+
+Example::
+
+        jt.concat([jt.array([[1],[2]]), jt.array([[2],[2]])], dim=1)
+        # return [[1],[2],[2],[2]]
+    '''
+    # TODO: low performance when concat lots of vars
+    total_dim = 0
+    if dim < 0: dim += len(arr[0].shape)
+    for a in arr:
+        total_dim += a.shape[dim]
+    cdim = 0
+    shape = list(a.shape)
+    shape[dim] = total_dim
+    s = jt.empty(shape, a.dtype)
+    slices = [slice(None)]*len(a.shape)
+    for a in arr:
+        if a.shape[dim] == 0:
+            continue
+        slices[dim] = slice(cdim, cdim+a.shape[dim])
+        # print(slices, type(a))
+        s = s.setitem(tuple(slices), a)
+        # s = jt.setitem(s, tuple(slices), a)
+        cdim += a.shape[dim]
+    return s

python/jittor/init.py

@@ -56,6 +56,45 @@ def relu_invariant_gauss(shape, dtype, mode="fan_in"):
 def relu_invariant_gauss_(var, mode="fan_in"):
     var.assign(relu_invariant_gauss(tuple(var.shape), var.dtype, mode))
 
+def calculate_std(var,mode,nonlinearity,param=0.01):
+    mode = mode.lower()
+    assert isinstance(param,(int,float))
+    assert var.ndim>=2
+    assert mode in ['fan_in', 'fan_out']
+
+    fan = var.shape[1] if mode == 'fan_in' else var.shape[0]
+    fan *= var[0][0].numel()
+
+    gains = {
+        'linear':1,
+        'conv1d':1,
+        'conv2d':1,
+        'conv3d':1,
+        'conv_transpose1d':1,
+        'conv_transpose2d':1,
+        'conv_transpose3d':1,
+        'sigmoid':1,
+        'tanh':5.0/3,
+        'relu':math.sqrt(2.0),
+        'leaky_relu':math.sqrt(2.0 / (1 + param ** 2)),
+    }
+    gain = gains[nonlinearity]
+    std = gain/math.sqrt(fan)
+    return std
+
+
+def kaiming_uniform_(var, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+    std = calculate_std(var,mode,nonlinearity,a)
+    bound = math.sqrt(3.0) * std 
+    with jt.no_grad():
+        return uniform_(var,-bound, bound)
+
+def kaiming_normal_(var, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+    std = calculate_std(var,mode,nonlinearity,a)
+    with jt.no_grad():
+        return gauss_(var,0, std)
+
+
 #TODO: bound = gain * math.sqrt(6.0/fan) ??
 def xavier_uniform(shape, dtype, gain=1.0):
     assert len(shape)>1
@@ -81,4 +120,4 @@ def xavier_gauss(shape, dtype, gain=1.0):
     return gauss(shape, dtype, 0, std)
 
 def xavier_gauss_(var, gain=1.0):
-    var.assign(xavier_gauss(tuple(var.shape), var.dtype, gain))
+    var.assign(xavier_gauss(tuple(var.shape), var.dtype, gain))