Jittor / jittor

Jittor is a high-performance deep learning framework based on JIT compiling and meta-operators.
https://cg.cs.tsinghua.edu.cn/jittor/
Apache License 2.0
3.08k stars 311 forks source link

某些原地操作函数传入slice无效 #367

Open LDYang694 opened 2 years ago

LDYang694 commented 2 years ago

Describe the bug

某些原地操作函数传入slice无效,如jt.init.uniform_(x[:,0]),可能是assign函数无法对slice进行赋值

Full Log

Minimal Reproduce

示例代码

import jittor as jt
a=jt.rand(3,3)
print("init before",a)
jt.init.uniform_(a[:,0])
print("init after",a)

import torch
a=torch.rand(3,3)
print("init before",a)
torch.nn.init.normal_(a[:,0])
print("init after",a)

运行结果

init before jt.Var([[0.18711373 0.8203873  0.24921839]
 [0.61351943 0.42086402 0.4616502 ]
 [0.95473033 0.15286785 0.25007996]], dtype=float32)
init after jt.Var([[0.18711373 0.8203873  0.24921839]
 [0.61351943 0.42086402 0.4616502 ]
 [0.95473033 0.15286785 0.25007996]], dtype=float32)
init before tensor([[0.8754, 0.2284, 0.9921],
        [0.4754, 0.0318, 0.5537],
        [0.2073, 0.6610, 0.1420]])
init after tensor([[-0.1734,  0.2284,  0.9921],
        [-1.1707,  0.0318,  0.5537],
        [ 0.5498,  0.6610,  0.1420]])

函数定义

def uniform_(var, low=0, high=1):
    ''' Inplace initialize Jittor Var by random uniform.
    '''
    return var.assign(uniform(var.shape, var.dtype, low, high))

Expected behavior

原地操作函数应该对于slice的支持?

diamond0910 commented 2 years ago

Thank you! We will fix this.