unsupported keyword argument "time" when taking a gradient with Zygote

Doing the following outputs a reasonable error? Any hints? Doing a similar function without AxisKeys works, however the advantage of using the time dimension is lost, which I would like to keep.

I suppose that I will need a new rule, however not sure how to actually start doing it.

using AxisKeys
using AxisKeys: KeyedArray as KA
using Zygote

ab = (;
    a =KA([5.0f0, 10.0f0];  time=1:2),
    b = KA([-2.0f0, 0.1f0];  time=1:2),
    )

function getVals(ab::NamedTuple, ts::Int)
    map(ab) do v
        in(:time, AxisKeys.dimnames(v)) ? v[time=ts][1] : v
    end
end

gradient(x -> x^2 + x*sum(getVals(ab,2)), 5)

ERROR: MethodError: no method matching adjoint(::Zygote.Context{false}, ::typeof(getindex), ::KeyedArray{Float32, 1, NamedDimsArray{(:time,), Float32, 1, Vector{Float32}}, Base.RefValue{UnitRange{Int64}}}; time=2)
Closest candidates are:
  adjoint(::ZygoteRules.AContext, ::typeof(getindex), ::AbstractArray, ::Any...) at none:0 got unsupported keyword argument "time"
  adjoint(::ZygoteRules.AContext, ::Base.Fix2, ::Any) at none:0 got unsupported keyword argument "time"
  adjoint(::ZygoteRules.AContext, ::Base.Fix1, ::Any) at none:0 got unsupported keyword argument "time"
  ...
Stacktrace:
  [1] _pullback
    @ ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:75 [inlined]
  [2] _pullback
.
.
.

No source file or test is named. Begin by running the provided Julia AxisKeys/Zygote reproduction, then trace the adjoint handling for keyword-based getindex with the time dimension. Done means the gradient example no longer raises the unsupported keyword argument error while retaining time-based indexing.