backward of Gather Args: dy (CTensor): gradient tensor. Returns: the gradient tensor over input tensor.
(self, dy)
| 4478 | return singa.ConcatOn(xs, self.axis) |
| 4479 | |
| 4480 | def backward(self, dy): |
| 4481 | """ |
| 4482 | backward of Gather |
| 4483 | Args: |
| 4484 | dy (CTensor): gradient tensor. |
| 4485 | Returns: |
| 4486 | the gradient tensor over input tensor. |
| 4487 | """ |
| 4488 | _shape = self.x_shape[self.axis] |
| 4489 | |
| 4490 | def construct_dx(dy, axis, indices, _shape): |
| 4491 | dys = [] |
| 4492 | data_idx = 0 |
| 4493 | data_idxes = tuple(indices) |
| 4494 | for step_idx in range(_shape): |
| 4495 | if step_idx in data_idxes: |
| 4496 | tmp_tensor = singa.SliceOn(dy, data_idx, data_idx + 1, axis) |
| 4497 | data_idx += 1 |
| 4498 | else: |
| 4499 | tmp_shape = list(dy.shape()) |
| 4500 | tmp_shape[axis] = 1 |
| 4501 | tmp_tensor = singa.Tensor(tmp_shape, dy.device()) |
| 4502 | tmp_tensor.SetFloatValue(0.) |
| 4503 | dys.append(tmp_tensor) |
| 4504 | dys = singa.VecTensor(dys) |
| 4505 | dy = singa.ConcatOn(dys, axis) |
| 4506 | return dy |
| 4507 | |
| 4508 | if isinstance(self.indices[0], tuple) or isinstance( |
| 4509 | self.indices[0], list): |
| 4510 | dx = singa.Tensor(self.x_shape, dy.device()) |
| 4511 | dx.SetFloatValue(0.) |
| 4512 | for data_idx in range(len(self.indices)): |
| 4513 | # get a piece of the dy and remove its new axis added at forward |
| 4514 | tmp_tensor = singa.SliceOn(dy, data_idx, data_idx + 1, |
| 4515 | self.axis) |
| 4516 | _slice_shape = list(tmp_tensor.shape()) |
| 4517 | del _slice_shape[self.axis] |
| 4518 | tmp_tensor = singa.Reshape(tmp_tensor, _slice_shape) |
| 4519 | # construct dx and sum them together |
| 4520 | tmp_tensor = construct_dx(tmp_tensor, self.axis, |
| 4521 | self.indices[data_idx], |
| 4522 | self.x_shape[self.axis]) |
| 4523 | dx = singa.__add__(dx, tmp_tensor) |
| 4524 | return dx |
| 4525 | else: |
| 4526 | return construct_dx(dy, self.axis, self.indices, _shape) |
| 4527 | |
| 4528 | |
| 4529 | def gather(x, axis, indices): |