Can I define a type-unstable EnzymeRule?

Specific Domains Machine Learning
,
1

I am trying to define custom Enzyme rules for my package ImplicitDifferentiation.jl, but I’m held back by type inference failures inside my rules. The typical error message looks like

Enzyme execution failed.
  Enzyme: incorrect return type of shadow-only forward custom rule -
BatchDuplicatedNoNeed{Vector{Float32}, 2} Type[Const{ImplicitFunction{var"#forward#44", var"#conditions#45", typeof(\), Nothing, Nothing}}, BatchDuplicated{Vector{Float32}, 2}] 
want just shadow type Tuple{Vector{Float32}, Vector{Float32}}
found Tuple{Any, Any}

Is there no way to define a custom rule if it is not fully inferrable?

2

The insides of your rule can be type unstable/inferrable, however the return presently should be the same (and inferrable) as the primal.

It’s been on our todo list to have Enzyme do the implicit a::Vector{Float32} on the inside so that rule writers don’t need to do the cast before the return, but we haven’t gotten to that yet. Contributions welcome!

In any case you should just be able to do return val::Tuple{Vector{Float32}, Vector{Float32}}

2 Likes