Gram matrix in Theano

gram_matrix.py

import theano as th
import theano.tensor as tt

def gaussian_kernel(x, y, sigma=1.):
    return tt.exp(-tt.sum(tt.square(x - y)) / sigma**2)
    
def gram_matrix(X, Y, kernel):
    M = X.shape[0]
    N = Y.shape[0]
    
    R, _ = th.scan(
        fn=lambda k: kernel(X[k // N], Y[k % N]),
        sequences=[tt.arange(M * N)])

    return R.reshape([M, N])

X = tt.dmatrix('X')
Y = tt.dmatrix('Y')
K = gram_matrix(X, Y, gaussian_kernel)

f = th.function([X, Y], [tt.sum(K), tt.grad(tt.sum(K), X)])