<filename>HyperSphere/GP/inference/inverse_bilinear_form.py
import torch
from torch.autograd import Function, Variable, gradcheck


class InverseBilinearForm(Function):

	@staticmethod
	def forward(ctx, vec_left, matrix, vec_right=None):
		if vec_right is None:
			vec_right = vec_left

		ctx.save_for_backward(vec_left, matrix, vec_right)

		jitter = 0
		eye_mat = torch.diag(matrix[:1, :1].clone().repeat(matrix.size(1), 1).view(-1) * 0 + 1)
		while True:
			try:
				vec_right_sol = torch.gesv(vec_right, matrix + eye_mat * jitter)[0]
				break
			except RuntimeError:
				jitter = torch.mean(torch.diag(matrix)) * 1e-6 if jitter == 0 else jitter * 10

		return torch.mm(vec_left.t(), vec_right_sol)

	@staticmethod
	def backward(ctx, grad_output):
		"""
		:param ctx: 
		:param grad_output: grad_output is assumed to be d[Scalar]/dK and size() is n1 x n2 
		:return: 
		"""
		vec_left, matrix, vec_right = ctx.saved_variables
		grad_vec_left = grad_matrix = grad_vec_right = None

		linear_solver = torch.gesv(torch.cat([vec_left, vec_right], 1), matrix)[0]
		vec_left_sol = linear_solver[:, :vec_left.size(1)]
		vec_right_sol = linear_solver[:, vec_left.size(1):]
		if ctx.needs_input_grad[0]:
			grad_vec_left = grad_output * vec_right_sol
		if ctx.needs_input_grad[1]:
			grad_matrix = grad_output * -torch.mm(vec_left_sol, vec_right_sol.t())
		if ctx.needs_input_grad[2]:
			grad_vec_right = grad_output * vec_left_sol

		return grad_vec_left, grad_matrix, grad_vec_right


if __name__ == '__main__':
	ndim = 5
	vec_left = Variable(torch.randn(ndim, 1), requires_grad=True)
	A = torch.randn(ndim, ndim)
	matrix = Variable(A.mm(A.t()) + torch.eye(ndim), requires_grad=True)
	vec_right = Variable(torch.randn(ndim, 1), requires_grad=True)

	# gradcheck doesn't have to pass all the time.
	test = gradcheck(InverseBilinearForm.apply, (vec_left, matrix, vec_right), eps=1e-4, atol=1e-3, rtol=1e-2)
	print(test)

<filename>HyperSphere/GP/inference/inverse_bilinear_form.py
import torch
from torch.autograd import Function, Variable, gradcheck


class InverseBilinearForm(Function):

	@staticmethod
	def forward(ctx, vec_left, matrix, vec_right=None):
		if vec_right is None:
			vec_right = vec_left

		ctx.save_for_backward(vec_left, matrix, vec_right)

		jitter = 0
		eye_mat = torch.diag(matrix[:1, :1].clone().repeat(matrix.size(1), 1).view(-1) * 0 + 1)
		while True:
			try:
				vec_right_sol = torch.gesv(vec_right, matrix + eye_mat * jitter)[0]
				break
			except RuntimeError:
				jitter = torch.mean(torch.diag(matrix)) * 1e-6 if jitter == 0 else jitter * 10

		return torch.mm(vec_left.t(), vec_right_sol)

	@staticmethod
	def backward(ctx, grad_output):
		"""
		:param ctx: 
		:param grad_output: grad_output is assumed to be d[Scalar]/dK and size() is n1 x n2 
		:return: 
		"""
		vec_left, matrix, vec_right = ctx.saved_variables
		grad_vec_left = grad_matrix = grad_vec_right = None

		linear_solver = torch.gesv(torch.cat([vec_left, vec_right], 1), matrix)[0]
		vec_left_sol = linear_solver[:, :vec_left.size(1)]
		vec_right_sol = linear_solver[:, vec_left.size(1):]
		if ctx.needs_input_grad[0]:
			grad_vec_left = grad_output * vec_right_sol
		if ctx.needs_input_grad[1]:
			grad_matrix = grad_output * -torch.mm(vec_left_sol, vec_right_sol.t())
		if ctx.needs_input_grad[2]:
			grad_vec_right = grad_output * vec_left_sol

		return grad_vec_left, grad_matrix, grad_vec_right


if __name__ == '__main__':
	ndim = 5
	vec_left = Variable(torch.randn(ndim, 1), requires_grad=True)
	A = torch.randn(ndim, ndim)
	matrix = Variable(A.mm(A.t()) + torch.eye(ndim), requires_grad=True)
	vec_right = Variable(torch.randn(ndim, 1), requires_grad=True)

	# gradcheck doesn't have to pass all the time.
	test = gradcheck(InverseBilinearForm.apply, (vec_left, matrix, vec_right), eps=1e-4, atol=1e-3, rtol=1e-2)
	print(test)

:param ctx:  		:param grad_output: grad_output is assumed to be d[Scalar]/dK and size() is n1 x n2  		:return: # gradcheck doesn't have to pass all the time.