Commit 7febeb31 authored by rusty1s's avatar rusty1s
Browse files

removed mul

parent 79239425
...@@ -6,16 +6,16 @@ void scatter_mul(at::Tensor src, at::Tensor index, at::Tensor out, ...@@ -6,16 +6,16 @@ void scatter_mul(at::Tensor src, at::Tensor index, at::Tensor out,
int64_t dim) { int64_t dim) {
int64_t elems_per_row = index.size(dim), i, idx; int64_t elems_per_row = index.size(dim), i, idx;
printf("elems_per_row: %lli\n", elems_per_row); printf("elems_per_row: %lli\n", elems_per_row);
AT_DISPATCH_ALL_TYPES(src.type(), "scatter_mul", [&] { // AT_DISPATCH_ALL_TYPES(src.type(), "scatter_mul", [&] {
DIM_APPLY3(scalar_t, src, int64_t, index, scalar_t, out, dim, { // DIM_APPLY3(scalar_t, src, int64_t, index, scalar_t, out, dim, {
for (i = 0; i < elems_per_row; i++) { // for (i = 0; i < elems_per_row; i++) {
idx = index_data[i * index_stride]; // idx = index_data[i * index_stride];
printf("i: %lli, idx: %lli\n", i, idx); // printf("i: %lli, idx: %lli\n", i, idx);
printf("src: %lli\n", (int64_t)src_data[i * src_stride]); // printf("src: %lli\n", (int64_t)src_data[i * src_stride]);
out_data[idx * out_stride] *= src_data[i * src_stride]; // out_data[idx * out_stride] *= src_data[i * src_stride];
} // }
}); // });
}); // });
} }
void scatter_div(at::Tensor src, at::Tensor index, at::Tensor out, void scatter_div(at::Tensor src, at::Tensor index, at::Tensor out,
......
...@@ -121,9 +121,10 @@ def test_forward(test, dtype, device): ...@@ -121,9 +121,10 @@ def test_forward(test, dtype, device):
op = getattr(torch_scatter, 'scatter_{}'.format(test['name'])) op = getattr(torch_scatter, 'scatter_{}'.format(test['name']))
out = op(src, index, test['dim'], fill_value=test['fill_value']) out = op(src, index, test['dim'], fill_value=test['fill_value'])
print(out)
if isinstance(out, tuple): # if isinstance(out, tuple):
assert out[0].tolist() == expected.tolist() # assert out[0].tolist() == expected.tolist()
assert out[1].tolist() == test['expected_arg'] # assert out[1].tolist() == test['expected_arg']
else: # else:
assert out.tolist() == expected.tolist() # assert out.tolist() == expected.tolist()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment