# Learning a mapping from an image to an ODE

**URL:** <https://discourse.julialang.org/t/learning-a-mapping-from-an-image-to-an-ode/134456>\
**Category:** Machine Learning\
**Tags:** ode, lux\
**Created:** [December 9, 2025, 11:11pm UTC](https://discourse.julialang.org/t/learning-a-mapping-from-an-image-to-an-ode/134456 "2025-12-09T23:11:01Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![csvance](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/csvance/32/218927_2.png) [@csvance](https://discourse.julialang.org/u/csvance)\
**Post date:** [December 9, 2025, 11:11pm UTC](https://discourse.julialang.org/t/learning-a-mapping-from-an-image-to-an-ode/134456/1 "2025-12-09T23:11:01Z")

</div>

I have trained a CNN to map an image to a sequence of feature vectors `f_v` and I want to learn a function `du(u, f_v)` that maps the feature vectors to an ODE. For each image I know how to compute both the function and first derivative that I am trying to model over the set of all images in my domain. My question is in most examples I see the input to the neural ODE is just `u0`. That’s fine for learning a single ODE, but from a practical standpoint during training how can I make my NeuralODE a function of both `u` and these features in Julia/Lux/OrdinaryDiffEQ? I don’t need to backprop through the features or optimize them as parameters, its a fixed representation per image.

More concretely say I want to use this Lux model for my NeuralODE:

```julia
model = @compact(
    Q=Dense(2 => 128),
    K=Dense(3 => 128),
    V=Dense(3 => 128),
    δU=Dense(128, 2),
    act=tanh
) do inps
    # u has shape (2,)
    # f has shape (3, 400)
    u, f_v = inps

    q = Q(u)
    k = K(f_v)
    v = V(f_v)

    attn = softmax(
        sum(q.*k, dims=1) ./ sqrt(128), 
        dims=2
    )

    # Apply the attention mask and reduce (128, 400) -> (128, 1)
    z = sum(attn.*v, dims=2)

    z = act.(z)

    du = δU(z)

    @return du
end

```

Is this something I can work with, or do I need to use some sort of trick to get the additional features as inputs? I have seen in other posts that additional parameters were passed with `u0`: [Using Neural ODEs to learn a family of ODEs (with Automatic Differentiation) - #14 by leespen1](https://discourse.julialang.org/t/using-neural-odes-to-learn-a-family-of-odes-with-automatic-differentiation/134348/14), but its not clear to me how that would impact optimization. Ultimately only `u[1]` and `u[2]` are what I want out of the ODE. I guess I could discount the loss on 1200 additional dimensions of `u`, but is this the most idiomatic/computationally efficient way to handle these types of cases? I’m completely new to doing ML in Julia and pretty fresh when it comes to neural ODE. I went over a decent amount of examples, but I’m still feeling pretty lost here.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 10, 2025, 1:13pm UTC](https://discourse.julialang.org/t/learning-a-mapping-from-an-image-to-an-ode/134456/2 "2025-12-10T13:13:40Z")

</div>

Just make the NN take in the vector `[u; f_v]` and make sure it’s sized correctly?

---

<div class="post-metadata">

**Author:** ![langestefan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/langestefan/32/207923_2.png) [@langestefan](https://discourse.julialang.org/u/langestefan)\
**Post date:** [December 10, 2025, 2:45pm UTC](https://discourse.julialang.org/t/learning-a-mapping-from-an-image-to-an-ode/134456/3 "2025-12-10T14:45:30Z")

</div>

Not sure how relevant it is for your usecase, but this came to mind: [Particle-filter tutorial · LowLevelParticleFilters Documentation](https://baggepinnen.github.io/LowLevelParticleFilters.jl/stable/beetle_example/#Smoothing-the-track-of-a-moving-beetle)

---

<div class="post-metadata">

**Author:** ![csvance](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/csvance/32/218927_2.png) [@csvance](https://discourse.julialang.org/u/csvance)\
**Post date:** [December 10, 2025, 5:47pm UTC](https://discourse.julialang.org/t/learning-a-mapping-from-an-image-to-an-ode/134456/4 "2025-12-10T17:47:37Z")

</div>

> [@ChrisRackauckas](#):
>
> Just make the NN take in the vector `[u; f_v]` and make sure it’s sized correctly?

I guess it’s not totally clear to me where these extra inputs belong. When working with OrdinaryDiffEQ.jl I could include them in my parameters type when I define the ODEProblem then access `p` inside my derivative function. Looking at the example in the Lux documentation here [MNIST Classification using Neural ODEs | Lux.jl Docs](https://lux.csail.mit.edu/stable/tutorials/intermediate/1_NeuralODE#Training) I see `Training.train_single_step!` providing an objective function and `(x, y)` as the data. What I don’t fully understand is the precise convention of `(x, y)` and how that is plugged into the model + loss function and solver.

Can my `x` just be `[u; f_v]`? If there are multiple inputs like this, is it just standard convention that the first one will be the one that gets integrated by the ODE solver?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 10, 2025, 6:00pm UTC](https://discourse.julialang.org/t/learning-a-mapping-from-an-image-to-an-ode/134456/5 "2025-12-10T18:00:06Z")

</div>

The the overload to:

```julia-auto
function (n::StatefulNeuralODE)(x, ps, st)
    st_model = StatefulLuxLayer(n.model, ps, st)
    dudt(u, p, t) = st_model([u; f_v] p)
    prob = ODEProblem{false}(ODEFunction{false}(dudt), x, n.tspan, ps)
    return solve(prob, n.solver; n.kwargs...), st_model.st
end

```

and you should be good if `f_v` is a cost function constant