not possible to train on GPU due to PDE

[armbrusl]

Hi,

I have recently run into an issue trying to train a bit more "complex" PDE on my gpu (it works fine but is quite slow on my cpu).

eqs = [
    (real(AB1(x))*Dxx(Eᵣ(x,z)) + real(AB2(x))*Dx(Eᵣ(x,z)) + real(AB3(z))*Dzz(Eᵣ(x,z)) + real(AB4(x))*Dz(Eᵣ(x,z)))/κ -
    (imag(AB1(x))*Dxx(Eᵢ(x,z)) + imag(AB2(x))*Dx(Eᵢ(x,z)) + imag(AB3(z))*Dzz(Eᵢ(x,z)) + imag(AB4(x))*Dz(Eᵢ(x,z)))/κ + κ*(ϵᵣ*Eᵣ(x,z) - ϵᵢ*Eᵢ(x,z)) ~ 0.0,
    (real(AB1(x))*Dxx(Eᵢ(x,z)) + real(AB2(x))*Dx(Eᵢ(x,z)) + real(AB3(z))*Dzz(Eᵢ(x,z)) + real(AB4(x))*Dz(Eᵢ(x,z)))/κ +
    (imag(AB1(x))*Dxx(Eᵣ(x,z)) + imag(AB2(x))*Dx(Eᵣ(x,z)) + imag(AB3(z))*Dzz(Eᵣ(x,z)) + imag(AB4(x))*Dz(Eᵣ(x,z)))/κ + κ*(ϵᵣ*Eᵢ(x,z) + ϵᵢ*Eᵣ(x,z)) ~ -J(z)
]

where the coefficients are this

function sigma(x, a, b)
    """
    sigma(x) = 0 if a < x < b, else grows cubically from zero
    """
    function _sigma(d)
        return σ₀ * d^2 * (1.0 / (1.0 + exp(-d/0.001)))
    end
    return _sigma(a - x) + _sigma(x - b)
end

function dsigma(x, a, b)
    function _sigma(d)
        return 2 * σ₀ * d * (1.0 / (1.0 + exp(-d/0.001)))
    end
    return -_sigma(a - x) + _sigma(x - b)
end

function AB1(x)
    sigma_x = sigma(x, xᵢ, xₛ)  
    AB1 = 1 / (1 + 1im / κ * sigma_x) ^ 2 
    return AB1
end

function AB2(x)
    sigma_x = sigma(x, xᵢ, xₛ)  
    dsigma_x = dsigma(x, xᵢ, xₛ)
    AB1 = 1 / (1 + 1im / κ * sigma_x) ^ 2 * dsigma_x
    AB2 = -1im / κ * AB1 / (1 + 1im / κ * sigma_x)
    return AB2
end

function AB3(z)
    sigma_z = sigma(z, zᵢ, zₛ)
    AB3 = 1 / (1 + 1im / κ * sigma_z) ^ 2
    return AB3
end

function AB4(z)
    sigma_z = sigma(z, zᵢ, zₛ)
    dsimga_z = dsigma(z, zᵢ, zₛ)
    AB3 = 1 / (1 + 1im / κ * sigma_z) ^ 2
    AB4 = -1im / κ * dsimga_z * AB3 / (1 + 1im / κ * sigma_z)
    return AB4
end


function J(z)
    return 1 / (0.2 * sqrt(pi)) * exp(-((z + 1.5)  / 0.2) ^ 2) 
end

The error (part of it) I get is

Stacktrace:
 [1] _getindex (repeats 85 times)
   @ ./broadcast.jl:669
 [2] _broadcast_getindex
   @ ./broadcast.jl:645
 [3] getindex
   @ ./broadcast.jl:605
 [4] #35
   @ ~/.julia/packages/GPUArrays/qt4ax/src/host/broadcast.jl:70
Reason: unsupported call to an unknown function (call to jl_f__apply_iterate)

The issue I think is with AB2 and AB4, if I remove them I can train on the gpu. I have also tried to simplify the functions to no avail.

Maybe somebody knows a quick fix for this or has some insight.

Thank you