# Creating a Neural ODE in Lux that explicitly depends on time

**URL:** <https://discourse.julialang.org/t/creating-a-neural-ode-in-lux-that-explicitly-depends-on-time/101062>\
**Category:** General Usage\
**Created:** [July 1, 2023, 8:12pm UTC](https://discourse.julialang.org/t/creating-a-neural-ode-in-lux-that-explicitly-depends-on-time/101062 "2023-07-01T20:12:24Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![George\_Stepaniants](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/george_stepaniants/32/47230_2.png) [@George\_Stepaniants](https://discourse.julialang.org/u/George_Stepaniants)\
**Post date:** [July 1, 2023, 8:12pm UTC](https://discourse.julialang.org/t/creating-a-neural-ode-in-lux-that-explicitly-depends-on-time/101062/1 "2023-07-01T20:12:24Z")

</div>

Hello,

I am very excited to use Lux for Neural ODE training but am having trouble understanding how to extend neural ODEs to accept time t as an argument.

I noticed this thread: [DiffEqFlux with time as additional input to Neural ODE - #3 by Michael\_Struwig](https://discourse.julialang.org/t/diffeqflux-with-time-as-additional-input-to-neural-ode/26456/3)  
which discusses how to do this for neural ODEs in Flux but this doesn’t seem to translate directly into the Lux interface. Any help would be much appreciated!

---

<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:** [July 2, 2023, 7:27am UTC](https://discourse.julialang.org/t/creating-a-neural-ode-in-lux-that-explicitly-depends-on-time/101062/2 "2023-07-02T07:27:37Z")

</div>

> [@George\_Stepaniants](#):
>
> I am very excited to use Lux for Neural ODE training but am having trouble understanding how to extend neural ODEs to accept time t as an argument.

Just put `t` into the neural network.

> [@George\_Stepaniants](#):
>
> I noticed this thread: [DiffEqFlux with time as additional input to Neural ODE - #3 by Michael\_Struwig](https://discourse.julialang.org/t/diffeqflux-with-time-as-additional-input-to-neural-ode/26456/3)  
> which discusses how to do this for neural ODEs in Flux but this doesn’t seem to translate directly into the Lux interface.

This thread’s main points all work exactly the same in the Lux interface. What did you try?

---

<div class="post-metadata">

**Author:** ![George\_Stepaniants](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/george_stepaniants/32/47230_2.png) [@George\_Stepaniants](https://discourse.julialang.org/u/George_Stepaniants)\
**Post date:** [July 2, 2023, 7:54pm UTC](https://discourse.julialang.org/t/creating-a-neural-ode-in-lux-that-explicitly-depends-on-time/101062/3 "2023-07-02T19:54:55Z")

</div>

I tried just putting “t” in the network but needed to actually redefine the neural ODE in the following way (below I allow this new neural ODE to accept any function of time “control(t)”)

```julia

function ControlledNeuralODE(model::Lux.AbstractExplicitLayer, control;
  solver=Tsit5(),
  sensealg=InterpolatingAdjoint(; autojacvec=ZygoteVJP()),
  tspan=(0.0f0, 1.0f0),
  kwargs...)
  return ControlledNeuralODE(model, control, solver, sensealg, tspan, kwargs)
end

function (n::ControlledNeuralODE)(x, ps, st)
  function dudt(u, p, t)
    c = n.control(t)
    u_, st = n.model([u; c], p, st)
    return u_
  end
  prob = ODEProblem(ODEFunction(dudt), x, n.tspan, ps)
  return solve(prob, n.solver; sensealg=n.sensealg, n.kwargs...), st
end

```

Just wanted to verify that this is a reasonable way of redefining the neural ODE. I followed the same outline as in: [MNIST Classification using NeuralODE - Lux.jl](https://lux.csail.mit.edu/stable/examples/generated/intermediate/NeuralODE/main/)

---

<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:** [July 2, 2023, 9:57pm UTC](https://discourse.julialang.org/t/creating-a-neural-ode-in-lux-that-explicitly-depends-on-time/101062/4 "2023-07-02T21:57:40Z")

</div>

That looks fine as long as the input size for the model is correct. Did it not work?

---

<div class="post-metadata">

**Author:** ![George\_Stepaniants](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/george_stepaniants/32/47230_2.png) [@George\_Stepaniants](https://discourse.julialang.org/u/George_Stepaniants)\
**Post date:** [July 2, 2023, 10:03pm UTC](https://discourse.julialang.org/t/creating-a-neural-ode-in-lux-that-explicitly-depends-on-time/101062/5 "2023-07-02T22:03:24Z")

</div>

Okay perfect, yes the way I wrote it seems to do the trick!