DiffEqGPU.jl with CUDA: Error computing gradients through SDE solver

Specific Domains Machine Learning
13

Unfortunately, that doesn’t seem to help. My code looks like this:

nu = 1.5
n_samples = 100
d = 4
l_list = ones((d,1))
samples_rff = rand(TDist(2*nu), (d, n_samples));
# Generate A from a normal distribution
A = randn((d, n_samples))/sqrt(n_samples)
# Send the parameters to the GPU
A = cu(A)
samples_rff = cu(samples_rff)
l_list = cu(l_list)
weights = hcat((A, l_list)...)

function rff_model(X, A, sample_features, l)
    tau = l.^(-1)
    W = tau.*sample_features
    W = sample_features
    M = W'*X
    M = cos.(M) + sin.(M)
    return A*M
end
forward_model(x, A, l) = rff_model(x, A, samples_rff, l)

# Define our drift and diffusion functions
function drift_rff(dstate, state, p, t)	
    A = p[:, 1:end-1]
    l = p[:, end:end]
    return forward_model(state, A, l)
end

function diffussion_rff(dstate,state,p,t)
# noise is a vector of size(4,1)
    return dstate .+ noise 
end

prob = SDEProblem{true}( drift_rff, diffussion_rff,  initial_conditions,(t_in, t_fin), weights)
function loss(weights, training_trajectories, initial_conditions, training_times)
    temp_prob = remake(prob, p = weights)
    tmp_sol = solve(temp_prob,SOSRA(),saveat=training_times,  force_dtmin = true);
    arrsol = CuArray(tmp_sol)
    return sum((arrsol - training_trajectories).^2) 
end
objective = weights -> loss(weights, training_trajectories, initial_conditions, training_times)
val, grads = Zygote.withgradient(objective, weights)

Is perhaps indexing the parameters like this

A = p[:, 1:end-1]
l = p[:, end:end]

forbidden? Or is using globally defined arguments to define my forward model a problem? My noise is a vector of size (4,1), could that be an issue? I’m somewhat confused because I feel like what I’m doing is fairly simple and computing the loss itself is not a problem.

Thank you for your help.

EDIT: after some more testing, I tried to define the problem out-of-place in the following way:

# Define our drift and diffusion functions
function drift_rff(state, p, t)	
    A = p[:, 1:n_samples]
    l = p[:, n_samples+1:n_samples+1]
    noise = p[:, end:end]
    return forward_model(state, A, l)
end
function diffussion_rff(state,p,t)
    A = p[:, 1:n_samples]
    l = p[:, n_samples+1:n_samples+1]
    noise = p[:, end:end]
    return noise
end

@show drift_rff(cu(u0), weights, 0.0)
@show diffussion_rff(cu(u0), weights, 0.0)
prob = SDEProblem{true}( drift_rff, diffussion_rff,  initial_conditions,(t_in, t_fin), weights)

but this throws me an error:

ERROR: LoadError: Nonconforming functions detected. If a model function f is defined as in-place, then all constituent functions like jac and paramjac must be in-place (and vice versa with out-of-place). Detected that some overloads did not conform to the same convention as f.

This error pops up no matter the definition of diffusion (even if it’s defined to be the drift itself!). I’m not sure what is going on. Note that my initial condition is either a vector of dimension d or a matrix of size dxn, so I’m solving a system of equations.