Custom stuct and Jacobian: Issue with Zygote not understanding custom struct

Specific Domains
,
1

The function with ‘normals’ input is fine but when it’s inside the struct, Zygote returns (nothing)

# MWE 
struct MeshExample
    normals::AbstractMatrix
    other::Float64
end
function radiation_using_struct(mesh::MeshExample, dof::AbstractVector, omega::AbstractFloat)
    return -1im * omega .* sum(mesh.normals .* dof', dims=2)
end

function radiation_using_normals(normals, dof::AbstractVector, omega::AbstractFloat)
    return -1im * omega .* sum(normals .* dof', dims=2)
end

normals = rand(10, 3)  
mesh = MeshExample(normals, 0.1)
dof = [0, 0, 1]  
omega = 2.0    

println(Zygote.jacobian(m->imag.(radiation_using_struct(mesh,dof,omega)),mesh))
println(Zygote.jacobian(m->imag.(radiation_using_normals(m,dof,omega)),normals))

Ideally, I would prefer to use mesh as input.

1 Like