Autodiff across FE result

[cpfiffer]

Someone noted to me that FixedEffectModels.jl is tricky to use AD on because there are so many explicit Float64 type constraints -- does anyone have a good sense of how much effort it would take to remove/parameterize/reduce the explicit type constraints here?

As an example, the FixedEffectModel struct has a lot of Float64 explicit types that could be parametric instead.

https://github.com/FixedEffects/FixedEffectModels.jl/blob/2a2661e508c579b6c9c8a82f1388debe0c8229e8/src/FixedEffectModel.jl#L8-L45

Is there an appetite for this? I think it'd be lovely to be able to AD through high-dimensional fixed effect estimates.

[schrimpf]

What do you have in mind? coef inherits its element type from y and X. These are constructed from the dataframe and formula, and converted to Float64. I could imagine getting the type of y and X from the types of the columns of the dataframe to allow automatic differentiation with respect to the data, but it seems somewhat unusual to try to autodiff through a dataframe.

If you use the lower level code in FixedEffects.jl, there is some compatibility with AD.

using ForwardDiff, FiniteDiff, FixedEffects, LinearAlgebra

N = 10
K = 2
d1=repeat(1:(N÷2), inner = 2)
d2=repeat(1:2, inner = N÷2)
p1 = FixedEffect(d1)
p2 = FixedEffect(d2)
x = rand(N,K)
y = rand(N)
function regfe(y,x,fes)
  x = copy(x)
  y = copy(y)
  w = FixedEffects.uweights(eltype(x), size(x, 1))
  feM = AbstractFixedEffectSolver{eltype(x)}(fes, w, Val{:cpu}, Threads.nthreads())
  solve_residuals!(x, feM)
  return((x'*x) \ (x'*y))
end
regfe(y, x, [p1, p2])


# derivative wrt x
J1 = ForwardDiff.jacobian(x->regfe(y, reshape(x,N,K), [p1, p2]) , vec(x))
J2 = FiniteDiff.finite_difference_jacobian(x->regfe(y, reshape(x,N,K), [p1, p2]) , vec(x))
norm(J1- J2)

# derivative wrt y
J1 = ForwardDiff.jacobian(y->regfe(y, reshape(x,N,K), [p1, p2]) , y)
J2 = FiniteDiff.finite_difference_jacobian(y->regfe(y, reshape(x,N,K), [p1, p2]) , y)
norm(J1- J2)

It's dreadfully slow, but appears to give correct results. Reverse mode packages are probably harder. The in-place mutation of solve_residuals! and iteration in lsmr probably mean it'd require defining custom methods. Doing so could also greatly improve the ForwardDiff performance.