def forward(ctx, X, Y, static_kernel, dyadic_order, sym=False, _naive_solver=False): A = X.shape[0] B = Y.shape[0] M = X.shape[1] N = Y.shape[1] D = X.shape[2] MM = (2**dyadic_order) * (M - 1) NN = (2**dyadic_order) * (N - 1) # computing dsdt k(X^i_s,Y^j_t) G_static = static_kernel.Gram_matrix(X, Y) G_static_ = G_static[:, :, 1:, 1:] + G_static[:, :, :-1, : -1] - G_static[:, :, 1:, : -1] - G_static[:, :, : -1, 1:] G_static_ = tile( tile(G_static_, 2, 2**dyadic_order) / float(2**dyadic_order), 3, 2 **dyadic_order) / float(2**dyadic_order) # if on GPU if X.device.type == 'cuda': assert max( MM, NN ) < 1024, 'n must be lowered or data must be moved to CPU as the current choice of n makes exceed the thread limit' # cuda parameters threads_per_block = max(MM + 1, NN + 1) n_anti_diagonals = 2 * threads_per_block - 1 # Prepare the tensor of output solutions to the PDE (forward) G = torch.zeros((A, B, MM + 2, NN + 2), device=G_static.device, dtype=G_static.dtype) G[:, :, 0, :] = 1. G[:, :, :, 0] = 1. # Run the CUDA kernel. blockspergrid = (A, B) compute_sig_kernel_Gram_mat_varpar_from_increments_cuda[ blockspergrid, threads_per_block](cuda.as_cuda_array(G_static_.detach()), MM + 1, NN + 1, n_anti_diagonals, cuda.as_cuda_array(G), _naive_solver) G = G[:, :, :-1, :-1] else: G = torch.tensor(sig_kernel_Gram_varpar(G_static_.detach().numpy(), sym, _naive_solver), dtype=G_static.dtype, device=G_static.device) ctx.save_for_backward(X, Y, G, G_static) ctx.sym = sym ctx.static_kernel = static_kernel ctx.dyadic_order = dyadic_order ctx._naive_solver = _naive_solver return G[:, :, -1, -1]
def backward(ctx, grad_output): X, Y, G, G_static = ctx.saved_tensors sym = ctx.sym static_kernel = ctx.static_kernel dyadic_order = ctx.dyadic_order _naive_solver = ctx._naive_solver G_static_ = G_static[:, :, 1:, 1:] + G_static[:, :, :-1, : -1] - G_static[:, :, 1:, : -1] - G_static[:, :, : -1, 1:] G_static_ = tile( tile(G_static_, 2, 2**dyadic_order) / float(2**dyadic_order), 3, 2 **dyadic_order) / float(2**dyadic_order) A = X.shape[0] B = Y.shape[0] M = X.shape[1] N = Y.shape[1] D = X.shape[2] MM = (2**dyadic_order) * (M - 1) NN = (2**dyadic_order) * (N - 1) # Reverse paths X_rev = torch.flip(X, dims=[1]) Y_rev = torch.flip(Y, dims=[1]) # computing dsdt k(X_rev^i_s,Y_rev^j_t) for variation of parameters G_static_rev = flip(flip(G_static_, dim=2), dim=3) # if on GPU if X.device.type == 'cuda': # Prepare the tensor of output solutions to the PDE (backward) G_rev = torch.zeros((A, B, MM + 2, NN + 2), device=G_static.device, dtype=G_static.dtype) G_rev[:, :, 0, :] = 1. G_rev[:, :, :, 0] = 1. # cuda parameters threads_per_block = max(MM + 1, NN + 1) n_anti_diagonals = 2 * threads_per_block - 1 # Compute signature kernel for reversed paths blockspergrid = (A, B) compute_sig_kernel_Gram_mat_varpar_from_increments_cuda[ blockspergrid, threads_per_block](cuda.as_cuda_array(G_static_rev.detach()), MM + 1, NN + 1, n_anti_diagonals, cuda.as_cuda_array(G_rev), _naive_solver) G_rev = G_rev[:, :, :-1, :-1] # if on CPU else: G_rev = torch.tensor(sig_kernel_Gram_varpar( G_static_rev.detach().numpy(), sym, _naive_solver), dtype=G_static.dtype, device=G_static.device) G_rev = flip(flip(G_rev, dim=2), dim=3) GG = G[:, :, :-1, :-1] * G_rev[:, :, 1:, 1:] # finite difference step h = 1e-9 Xh = X[:, :, :, None] + h * torch.eye( D, dtype=X.dtype, device=X.device)[None, None, :] Xh = Xh.permute(0, 1, 3, 2) Xh = Xh.reshape(A, M * D, D) G_h = static_kernel.Gram_matrix(Xh, Y) G_h = G_h.reshape(A, B, M, D, N) G_h = G_h.permute(0, 1, 2, 4, 3) Diff_1 = G_h[:, :, 1:, 1:, :] - G_h[:, :, 1:, :-1, :] - ( G_static[:, :, 1:, 1:])[:, :, :, :, None] + ( G_static[:, :, 1:, :-1])[:, :, :, :, None] Diff_1 = tile( tile(Diff_1, 2, 2**dyadic_order) / float(2**dyadic_order), 3, 2** dyadic_order) / float(2**dyadic_order) Diff_2 = G_h[:, :, 1:, 1:, :] - G_h[:, :, 1:, :-1, :] - ( G_static[:, :, 1:, 1:])[:, :, :, :, None] + ( G_static[:, :, 1:, :-1])[:, :, :, :, None] Diff_2 += -G_h[:, :, :-1, 1:, :] + G_h[:, :, :-1, :-1, :] + ( G_static[:, :, :-1, 1:])[:, :, :, :, None] - ( G_static[:, :, :-1, :-1])[:, :, :, :, None] Diff_2 = tile( tile(Diff_2, 2, 2**dyadic_order) / float(2**dyadic_order), 3, 2** dyadic_order) / float(2**dyadic_order) grad_1 = (GG[:, :, :, :, None] * Diff_1) / h grad_2 = (GG[:, :, :, :, None] * Diff_2) / h grad_1 = torch.sum(grad_1, axis=3) grad_1 = torch.sum(grad_1.reshape(A, B, M - 1, 2**dyadic_order, D), axis=3) grad_2 = torch.sum(grad_2, axis=3) grad_2 = torch.sum(grad_2.reshape(A, B, M - 1, 2**dyadic_order, D), axis=3) grad_prev = grad_1[:, :, :-1, :] + grad_2[:, :, 1:, :] # /¯¯ grad_next = torch.cat([ torch.zeros((A, B, 1, D), dtype=X.dtype, device=X.device), grad_1[:, :, 1:, :] ], dim=2) # / grad_incr = grad_prev - grad_1[:, :, 1:, :] grad_points = torch.cat( [(grad_2[:, :, 0, :] - grad_1[:, :, 0, :])[:, :, None, :], grad_incr, grad_1[:, :, -1, :][:, :, None, :]], dim=2) if sym: grad = (grad_output[:, :, None, None] * grad_points + grad_output.t()[:, :, None, None] * grad_points).sum(dim=1) return grad, None, None, None, None, None else: grad = (grad_output[:, :, None, None] * grad_points).sum(dim=1) return grad, None, None, None, None, None
def forward(ctx, K_XX, K_XY, K_YY, static_kernel, dyadic_order, lambda_, sym=False, _naive_solver=False, inspect=False, centered=False): A = K_XX.shape[0] B = K_YY.shape[0] M = K_XX.shape[2] N = K_YY.shape[2] MM = (2**dyadic_order) * (M - 1) NN = (2**dyadic_order) * (N - 1) # computing dsdt k(X^1[i]_s,Y^1[j]_t) G_base = innerprodCKME( K_XX, K_XY, K_YY, lambda_, static_kernel, sym=sym, centered=centered ) # <--------------------- this is the only change compared to rank 0 G_base_ = G_base[:, :, 1:, 1:] + G_base[:, :, :-1, : -1] - G_base[:, :, 1:, :-1] - G_base[:, :, :-1, 1:] G_base_ = tile( tile(G_base_, 2, 2**dyadic_order) / float(2**dyadic_order), 3, 2** dyadic_order) / float(2**dyadic_order) # if on GPU if K_XX.device.type == 'cuda': assert max( MM, NN ) < 1024, 'n must be lowered or data must be moved to CPU as the current choice of n makes exceed the thread limit' # cuda parameters threads_per_block = max(MM + 1, NN + 1) n_anti_diagonals = 2 * threads_per_block - 1 # Prepare the tensor of output solutions to the PDE (forward) G = torch.zeros((A, B, MM + 2, NN + 2), device=G_base.device, dtype=G_base.dtype) G[:, :, 0, :] = 1. G[:, :, :, 0] = 1. # Run the CUDA kernel. blockspergrid = (A, B) compute_sig_kernel_Gram_mat_varpar_from_increments_cuda[ blockspergrid, threads_per_block](cuda.as_cuda_array(G_base_.detach()), MM + 1, NN + 1, n_anti_diagonals, cuda.as_cuda_array(G), _naive_solver) G = G[:, :, :-1, :-1] else: G = torch.tensor(sig_kernel_Gram_varpar(G_base_.detach().numpy(), sym, _naive_solver), dtype=G_base.dtype, device=G_base.device) if inspect: return G[:, :, -1, -1], G_base return G[:, :, -1, -1]