| 5547 | self.condition = condition |
| 5548 | |
| 5549 | def forward(self, a, b): |
| 5550 | if isinstance(self.condition, list): |
| 5551 | self.condition = np.array(self.condition) |
| 5552 | if isinstance(self.condition, np.ndarray): |
| 5553 | self.condition = self.condition.astype(np.float32) |
| 5554 | self.condition = tensor.from_numpy(self.condition) |
| 5555 | self.condition.to_device(a.device()) |
| 5556 | self.condition = self.condition.data |
| 5557 | self.neg_condition = singa.AddFloat( |
| 5558 | singa.MultFloat(self.condition, -1.), 1.) |
| 5559 | _a, _b = a, b |
| 5560 | dtype0 = _a.data_type() |
| 5561 | dtype1 = _b.data_type() |
| 5562 | if dtype0 == singa.kInt or dtype1 == singa.kInt: |
| 5563 | _a = a.AsType(singa.kFloat32) |
| 5564 | _b = b.AsType(singa.kFloat32) |
| 5565 | res = singa.__add__(singa.__mul__(self.condition, _a), |
| 5566 | singa.__mul__(self.neg_condition, _b)) |
| 5567 | res = res.AsType(singa.kInt) |
| 5568 | else: |
| 5569 | res = singa.__add__(singa.__mul__(self.condition, _a), |
| 5570 | singa.__mul__(self.neg_condition, _b)) |
| 5571 | return res |
| 5572 | |
| 5573 | def backward(self, dy): |
| 5574 | da = singa.__mul__(self.condition, dy) |