| import torch |
| |
| from ..function import Function, InplaceFunction |
| |
| |
| # TODO: no need to save all args if the grad w.r.t. some of them is not needed |
| class _BlasBase(InplaceFunction): |
| |
| def __init__(self, alpha=1, beta=1, inplace=False): |
| super(_BlasBase, self).__init__(inplace) |
| self.alpha = alpha |
| self.beta = beta |
| |
| def _get_output(self, arg): |
| if self.inplace: |
| self.mark_dirty(arg) |
| return arg |
| else: |
| return arg.new().resize_as_(arg) |
| |
| |
| class Addmm(_BlasBase): |
| |
| def forward(self, add_matrix, matrix1, matrix2): |
| self.save_for_backward(matrix1, matrix2) |
| output = self._get_output(add_matrix) |
| return torch.addmm(self.alpha, add_matrix, self.beta, |
| matrix1, matrix2, out=output) |
| |
| def backward(self, grad_output): |
| matrix1, matrix2 = self.saved_tensors |
| grad_add_matrix = grad_matrix1 = grad_matrix2 = None |
| |
| if self.needs_input_grad[0]: |
| grad_add_matrix = grad_output |
| if self.alpha != 1: |
| grad_add_matrix = grad_add_matrix.mul(self.alpha) |
| |
| if self.needs_input_grad[1]: |
| grad_matrix1 = torch.mm(grad_output, matrix2.t()) |
| if self.beta != 1: |
| grad_matrix1 *= self.beta |
| |
| if self.needs_input_grad[2]: |
| grad_matrix2 = torch.mm(matrix1.t(), grad_output) |
| if self.beta != 1: |
| grad_matrix2 *= self.beta |
| |
| return grad_add_matrix, grad_matrix1, grad_matrix2 |
| |
| |
| class Addbmm(_BlasBase): |
| |
| def forward(self, add_matrix, batch1, batch2): |
| self.save_for_backward(batch1, batch2) |
| output = self._get_output(add_matrix) |
| return torch.addbmm(self.alpha, add_matrix, self.beta, |
| batch1, batch2, out=output) |
| |
| def backward(self, grad_output): |
| batch1, batch2 = self.saved_tensors |
| grad_add_matrix = grad_batch1 = grad_batch2 = None |
| |
| if self.needs_input_grad[0]: |
| grad_add_matrix = grad_output |
| if self.alpha != 1: |
| grad_add_matrix = grad_add_matrix.mul(self.alpha) |
| |
| if any(self.needs_input_grad[1:]): |
| batch_grad_output = (grad_output |
| .unsqueeze(0) |
| .expand(batch1.size(0), batch1.size(1), batch2.size(2))) |
| |
| if self.needs_input_grad[1]: |
| grad_batch1 = torch.bmm(batch_grad_output, batch2.transpose(1, 2)) |
| if self.beta != 1: |
| grad_batch1 *= self.beta |
| |
| if self.needs_input_grad[2]: |
| grad_batch2 = torch.bmm(batch1.transpose(1, 2), batch_grad_output) |
| if self.beta != 1: |
| grad_batch2 *= self.beta |
| |
| return grad_add_matrix, grad_batch1, grad_batch2 |
| |
| |
| class Baddbmm(_BlasBase): |
| |
| def forward(self, add_batch, batch1, batch2): |
| self.save_for_backward(batch1, batch2) |
| output = self._get_output(add_batch) |
| return torch.baddbmm(self.alpha, add_batch, self.beta, |
| batch1, batch2, out=output) |
| |
| def backward(self, grad_output): |
| batch1, batch2 = self.saved_tensors |
| grad_add_batch = grad_batch1 = grad_batch2 = None |
| |
| if self.needs_input_grad[0]: |
| grad_add_batch = grad_output |
| if self.alpha != 1: |
| grad_add_batch = grad_add_batch.mul(self.alpha) |
| |
| if self.needs_input_grad[1]: |
| grad_batch1 = torch.bmm(grad_output, batch2.transpose(1, 2)) |
| if self.beta != 1: |
| grad_batch1 *= self.beta |
| |
| if self.needs_input_grad[2]: |
| grad_batch2 = torch.bmm(batch1.transpose(1, 2), grad_output) |
| if self.beta != 1: |
| grad_batch2 *= self.beta |
| |
| return grad_add_batch, grad_batch1, grad_batch2 |
| |
| |
| class Addmv(_BlasBase): |
| |
| def forward(self, add_vector, matrix, vector): |
| self.save_for_backward(matrix, vector) |
| output = self._get_output(add_vector) |
| return torch.addmv(self.alpha, add_vector, self.beta, |
| matrix, vector, out=output) |
| |
| def backward(self, grad_output): |
| matrix, vector = self.saved_tensors |
| grad_add_vector = grad_matrix = grad_vector = None |
| |
| if self.needs_input_grad[0]: |
| grad_add_vector = grad_output |
| if self.alpha != 1: |
| grad_add_vector = grad_add_vector.mul(self.alpha) |
| |
| if self.needs_input_grad[1]: |
| grad_matrix = torch.ger(grad_output, vector) |
| if self.beta != 1: |
| grad_matrix *= self.beta |
| |
| if self.needs_input_grad[2]: |
| grad_vector = torch.mv(matrix.t(), grad_output) |
| if self.beta != 1: |
| grad_vector *= self.beta |
| |
| return grad_add_vector, grad_matrix, grad_vector |
| |
| |
| class Addr(_BlasBase): |
| |
| def forward(self, add_matrix, vector1, vector2): |
| self.save_for_backward(vector1, vector2) |
| output = self._get_output(add_matrix) |
| return torch.addr(self.alpha, add_matrix, self.beta, |
| vector1, vector2, out=output) |
| |
| def backward(self, grad_output): |
| vector1, vector2 = self.saved_tensors |
| grad_add_matrix = grad_vector1 = grad_vector2 = None |
| |
| if self.needs_input_grad[0]: |
| grad_add_matrix = grad_output |
| if self.alpha != 1: |
| grad_add_matrix = grad_add_matrix.mul(self.alpha) |
| |
| if self.needs_input_grad[1]: |
| grad_vector1 = torch.mv(grad_output, vector2) |
| if self.beta != 1: |
| grad_vector1 *= self.beta |
| |
| if self.needs_input_grad[2]: |
| # TODO: maybe it's better to do transpose + mv + transpose |
| grad_vector2 = torch.mm(vector1.unsqueeze(0), grad_output).squeeze(0) |
| if self.beta != 1: |
| grad_vector2 *= self.beta |
| |
| return grad_add_matrix, grad_vector1, grad_vector2 |
| |
| |
| class Dot(Function): |
| |
| def forward(self, vector1, vector2): |
| self.save_for_backward(vector1, vector2) |
| self.sizes = (vector1.size(), vector2.size()) |
| return vector1.new((vector1.dot(vector2),)) |
| |
| def backward(self, grad_output): |
| vector1, vector2 = self.saved_tensors |
| grad_vector1 = grad_vector2 = None |
| |
| if self.needs_input_grad[0]: |
| grad_vector1 = vector2.mul(grad_output[0]).view(self.sizes[0]) |
| |
| if self.needs_input_grad[1]: |
| grad_vector2 = vector1.mul(grad_output[0]).view(self.sizes[1]) |
| |
| return grad_vector1, grad_vector2 |
