Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix gather kernel 0 shape #5888

Merged
merged 6 commits into from
Aug 15, 2021
Merged

fix gather kernel 0 shape #5888

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
813b4b5
fix gather kernel 0 shape
simonJJJ Aug 14, 2021
c33e74a
Merge branch 'master' of https://github.com/Oneflow-Inc/oneflow into …
simonJJJ Aug 15, 2021
815d286
Merge branch 'master' into fix_gather_0shape
oneflow-ci-bot Aug 15, 2021
c7c34a6
recover module test
simonJJJ Aug 15, 2021
bf91fd3
Merge branch 'fix_gather_0shape' of https://github.com/Oneflow-Inc/on…
simonJJJ Aug 15, 2021
8eb191d
Merge branch 'master' into fix_gather_0shape
oneflow-ci-bot Aug 15, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions oneflow/user/kernels/gather_kernel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,7 @@ class GatherKernel final : public user_op::OpKernel, public user_op::CudaGraphSu
const int64_t axis = ctx->Attr<int64_t>("axis");
const int64_t num_indices = indices->shape().elem_cnt();
user_op::Tensor* out = ctx->Tensor4ArgNameAndIndex("out", 0);
if (out->shape().elem_cnt() == 0) { return; }

int64_t offset = 0;
if (state != nullptr) {
Expand Down
166 changes: 83 additions & 83 deletions python/oneflow/test/modules/test_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,89 +32,89 @@ def np_relu(np_arr):

@flow.unittest.skip_unless_1n1d()
class TestModule(flow.unittest.TestCase):
# def test_nested_module(test_case):
# class CustomModule(flow.nn.Module):
# def __init__(self):
# super().__init__()
# self.relu = flow.nn.ReLU()

# def forward(self, x):
# return self.relu(x)

# m = CustomModule()
# x = flow.Tensor(2, 3)
# flow.nn.init.uniform_(x, a=-1.0, b=1.0)
# y = m(x)
# test_case.assertTrue(np.array_equal(np_relu(x.numpy()), y.numpy()))

# def test_relu(test_case):
# relu = flow.nn.ReLU()
# x = flow.Tensor(2, 3)
# flow.nn.init.uniform_(x, a=-1.0, b=1.0)
# y = relu(x)
# test_case.assertTrue(np.array_equal(np_relu(x.numpy()), y.numpy()))

# def test_load_state_dict(test_case):
# class CustomModule(flow.nn.Module):
# def __init__(self):
# super().__init__()
# self.w = flow.nn.Parameter(flow.Tensor(2, 3))

# def forward(self, x):
# return self.w

# m = CustomModule()
# ones = np.ones((2, 3), dtype=np.float32)
# m.load_state_dict({"w": ones})
# x = flow.Tensor(2, 3)
# y = m(x).numpy()
# test_case.assertTrue(np.array_equal(y, ones))

# def test_state_dict(test_case):
# class CustomModule(flow.nn.Module):
# def __init__(self, param1, param2):
# super().__init__()
# self.param1 = param1
# self.param2 = param2

# tensor0 = flow.nn.Parameter(flow.Tensor(2, 3))
# tensor1 = flow.nn.Parameter(flow.Tensor(2, 3))
# sub_module = CustomModule(tensor0, tensor1)
# m = CustomModule(tensor1, sub_module)
# state_dict = m.state_dict()
# test_case.assertEqual(
# state_dict,
# {"param2.param1": tensor0, "param2.param2": tensor1, "param1": tensor1},
# )

# def test_parameter(test_case):
# shape = (3, 4)
# t = flow.Tensor(*shape)
# p = flow.nn.Parameter(t)
# test_case.assertEqual(type(p), flow.nn.Parameter)
# test_case.assertEqual(p.shape, shape)

# def test_module_forward(test_case):
# class CustomModule(flow.nn.Module):
# def __init__(self, w):
# super().__init__()
# self.w = w

# def forward(self, x):
# return x + self.w

# m = CustomModule(5)
# test_case.assertEqual(m(1), 6)
# m = CustomModule(4)
# test_case.assertEqual(m(3), 7)

# def test_train_eval(test_case):
# m = flow.nn.Module()
# test_case.assertEqual(m.training, True)
# m.train()
# test_case.assertEqual(m.training, True)
# m.eval()
# test_case.assertEqual(m.training, False)
def test_nested_module(test_case):
class CustomModule(flow.nn.Module):
def __init__(self):
super().__init__()
self.relu = flow.nn.ReLU()

def forward(self, x):
return self.relu(x)

m = CustomModule()
x = flow.Tensor(2, 3)
flow.nn.init.uniform_(x, a=-1.0, b=1.0)
y = m(x)
test_case.assertTrue(np.array_equal(np_relu(x.numpy()), y.numpy()))

def test_relu(test_case):
relu = flow.nn.ReLU()
x = flow.Tensor(2, 3)
flow.nn.init.uniform_(x, a=-1.0, b=1.0)
y = relu(x)
test_case.assertTrue(np.array_equal(np_relu(x.numpy()), y.numpy()))

def test_load_state_dict(test_case):
class CustomModule(flow.nn.Module):
def __init__(self):
super().__init__()
self.w = flow.nn.Parameter(flow.Tensor(2, 3))

def forward(self, x):
return self.w

m = CustomModule()
ones = np.ones((2, 3), dtype=np.float32)
m.load_state_dict({"w": ones})
x = flow.Tensor(2, 3)
y = m(x).numpy()
test_case.assertTrue(np.array_equal(y, ones))

def test_state_dict(test_case):
class CustomModule(flow.nn.Module):
def __init__(self, param1, param2):
super().__init__()
self.param1 = param1
self.param2 = param2

tensor0 = flow.nn.Parameter(flow.Tensor(2, 3))
tensor1 = flow.nn.Parameter(flow.Tensor(2, 3))
sub_module = CustomModule(tensor0, tensor1)
m = CustomModule(tensor1, sub_module)
state_dict = m.state_dict()
test_case.assertEqual(
state_dict,
{"param2.param1": tensor0, "param2.param2": tensor1, "param1": tensor1},
)

def test_parameter(test_case):
shape = (3, 4)
t = flow.Tensor(*shape)
p = flow.nn.Parameter(t)
test_case.assertEqual(type(p), flow.nn.Parameter)
test_case.assertEqual(p.shape, shape)

def test_module_forward(test_case):
class CustomModule(flow.nn.Module):
def __init__(self, w):
super().__init__()
self.w = w

def forward(self, x):
return x + self.w

m = CustomModule(5)
test_case.assertEqual(m(1), 6)
m = CustomModule(4)
test_case.assertEqual(m(3), 7)

def test_train_eval(test_case):
m = flow.nn.Module()
test_case.assertEqual(m.training, True)
m.train()
test_case.assertEqual(m.training, True)
m.eval()
test_case.assertEqual(m.training, False)

def test_module_setattr(test_case):
class CustomModule(flow.nn.Module):
Expand Down