| 1486 | self.axis = axis |
| 1487 | |
| 1488 | def forward(self, x): |
| 1489 | x_shape = x.shape() |
| 1490 | x_rank = len(x_shape) |
| 1491 | if isinstance(self.indices, Tensor): |
| 1492 | self.indices = tensor.to_numpy(self.indices) |
| 1493 | elif isinstance(self.indices, (list, tuple)): |
| 1494 | self.indices = np.array(self.indices) |
| 1495 | if isinstance(self.updates, Tensor): |
| 1496 | self.updates = tensor.to_numpy(self.updates) |
| 1497 | elif isinstance(self.updates, (list, tuple)): |
| 1498 | self.updates = np.array(self.updates) |
| 1499 | self.updates.astype(np.int32) |
| 1500 | _x = tensor.to_numpy(tensor.from_raw_tensor(x)) |
| 1501 | _x = _x.astype(np.float32) |
| 1502 | |
| 1503 | assert x_rank == 2, "Only support 2D input." |
| 1504 | assert x_rank == len( |
| 1505 | self.indices.shape |
| 1506 | ), "Index should have the same number of dimensions as output" |
| 1507 | assert -x_rank < self.axis <= x_rank, "Axis is out of range" |
| 1508 | assert np.logical_and( |
| 1509 | -_x.shape[self.axis] < self.indices, |
| 1510 | self.indices <= _x.shape[self.axis]).all( |
| 1511 | ), "The values of the indexes should be between %d and %d" % ( |
| 1512 | -_x.shape[self.axis], _x.shape[self.axis] - 1) |
| 1513 | |
| 1514 | self.axis = self.axis % x_rank |
| 1515 | u_shape = self.updates.shape |
| 1516 | y = _x.copy() |
| 1517 | for i in range(u_shape[0]): |
| 1518 | for j in range(u_shape[1]): |
| 1519 | idx = int(self.indices[i][j]) |
| 1520 | if self.axis == 0: |
| 1521 | y[idx][j] = self.updates[i][j] |
| 1522 | else: |
| 1523 | y[i][idx] = self.updates[i][j] |
| 1524 | y = tensor.from_numpy(y) |
| 1525 | y.to_device(x.device()) |
| 1526 | return y.data |
| 1527 | |
| 1528 | def backward(self, dy): |
| 1529 | mask = np.ones(dy.shape(), dtype=np.float32) |