Parameter Gradient JVP with Enzyme [Fix with destructure?]

[aksuhton]

I'm encountering an issue while moving my JVP and VJP code from Zygote/ForwardDiff to Enzyme. The docs suggest that what I'm doing is supported, but it looks like the tests in LuxDL only cover the input gradient case. A minimal working example is below, containing a potential fix, followed by the error message.

To achieve the fix I make use of destructure and calling Enzyme.autodiff directly. I'm happy to work on integrating this structure into Lux if desired.

      p, re = destructure(ps)
      f = Lux.@closure((x, p) -> sm(x, re(p)))
      jvp = only(@jit Enzyme.autodiff(
              ad.mode, Lux.Utils.annotate_enzyme_function(ad, f), Const(input), Duplicated(p, deepcopy(p))
      ))

using Lux
using Optimisers
using Reactant, Enzyme
using Random
##
  function JVP()
      ##
      seed = 0
      rng = Xoshiro(seed)
      dev = reactant_device()

      ## Pass
      model = Dense(5 => 4; use_bias = false)
      ps, st = Lux.setup(rng, model)
      ps = ps |> dev
      st = st |> dev
      input = rand(rng, Float32, 5, 2) |> dev
      target = rand(rng, Float32, 4, 2) |> dev
      sm = StatefulLuxLayer{true}(model, deepcopy(ps), st)
      sm_input = Base.Fix1(sm, deepcopy(input))
      jvp = @jit jacobian_vector_product(
          sm_input,
          AutoEnzyme(),
          ps,
          ps,
      )

      ## Fail
      model = Dense(5 => 4)
      ps, st = Lux.setup(rng, model)
      ps = ps |> dev
      st = st |> dev
      input = rand(rng, Float32, 5, 2) |> dev
      target = rand(rng, Float32, 4, 2) |> dev
      sm = StatefulLuxLayer{true}(model, deepcopy(ps), st)
      sm_input = Base.Fix1(sm, deepcopy(input))
      jvp = @jit jacobian_vector_product(
          sm_input,
          AutoEnzyme(),
          ps,
          ps,
      )

      ## Fix
      model = Chain(Dense(5 => 4, swish), Dense(4 => 4))
      ps, st = Lux.setup(rng, model)
      ps = ps |> dev
      st = st |> dev
      input = rand(rng, Float32, 5, 2) |> dev
      target = rand(rng, Float32, 4, 2) |> dev
      sm = StatefulLuxLayer{true}(model, deepcopy(ps), st)
      sm_input = Base.Fix1(sm, deepcopy(input))

      ad = Lux.Utils.normalize_autoenzyme_mode(Forward, AutoEnzyme())
      p, re = destructure(ps)
      f = Lux.@closure((x, p) -> sm(x, re(p)))
      jvp = only(@jit Enzyme.autodiff(
              ad.mode, Lux.Utils.annotate_enzyme_function(ad, f), Const(input), Duplicated(p, deepcopy(p))
      ))
      ##
      return jvp
  end

error: type of return operand 3 ('tensor<4x4xf32>') doesn't match function result type ('tensor<4xf32>') in function @jacobian_vector_product
┌ Error: Compilation failed, MLIR module written to /var/folders/mn/z3y89snn7tgggjn3d522c4v40000gs/T/reactant_SPhOV7/module_000_nxc9_post_all_pm.mlir
└ @ Reactant.MLIR.IR ~/.julia/packages/Reactant/LiDEK/src/mlir/IR/Pass.jl:146
ERROR: LoadError: "failed to run pass manager on module"
Stacktrace:
  [1] run!(pm::Reactant.MLIR.IR.PassManager, operation::Reactant.MLIR.IR.Operation, key::String)
    @ Reactant.MLIR.IR ~/.julia/packages/Reactant/LiDEK/src/mlir/IR/Pass.jl:194
  [2] run_pass_pipeline!(mod::Reactant.MLIR.IR.Module, pass_pipeline::String, key::String; enable_verifier::Bool)
    @ Reactant.Compiler ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:1486
  [3] run_pass_pipeline!
    @ ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:1481 [inlined]
  [4] compile_mlir!(mod::Reactant.MLIR.IR.Module, f::typeof(jacobian_vector_product), args::Tuple{…}, compile_options::CompileOptions, debugcache::Vector{…}
, callcache::Dict{…}, sdycache::Dict{…}, sdygroupidcache::Tuple{…}; fn_kwargs::@NamedTuple{}, backend::String, runtime::Val{…}, legalize_stablehlo_to_mhlo::
Bool, client::Reactant.XLA.PJRT.Client, kwargs::@Kwargs{})
    @ Reactant.Compiler ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:2065
  [5] compile_xla(ctx::Reactant.MLIR.IR.Context, f::Function, args::Tuple{…}; before_xla_optimizations::Bool, client::Nothing, serializable::Bool, kwargs::@
Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:4030
  [6] compile_xla
    @ ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:4030 [inlined]
  [7] compile(ctx::Reactant.MLIR.IR.Context, f::Function, args::Tuple{…}; kwargs::@Kwargs{…})
    @ Reactant.Compiler ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:4150
  [8] macro expansion
    @ ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:3093 [inlined]
  [9] macro expansion
    @ ~/.julia/packages/LLVM/fEIbx/src/base.jl:97 [inlined]
 [10] macro expansion
    @ ~/.julia/packages/Reactant/LiDEK/src/Compiler.jl:3092 [inlined]
 [11] JVP()
    @ Main ~/GitLab/Archon/mwe.jl:67
 [12] top-level scope
    @ ~/GitLab/Archon/mwe.jl:83
 [13] include(mapexpr::Function, mod::Module, _path::String)
    @ Base ./Base.jl:307
 [14] top-level scope

[avik-pal]

Can you open a corresponding issue in Reactant? This is a pass failure issue that needs to be handled upstream

[aksuhton]

The corresponding reactant issue is below. Thank you!

EnzymeAD/Reactant.jl#2734