NeuralODE implementation using JAX

[Sudhanshu-G-Ogalapur]

No description provided.

[lockwo]

1. format code as one block by wrapping it in "```" see raw of

print("Hello World")

2. Provide MVC (https://stackoverflow.com/help/minimal-reproducible-example), I don't have those files and thus can't run the code.

[lockwo]

MVC is the minimal code snippet that you can just copy into editor and have it work (not having to download files, which also presents a security risk). However, without I running the code I would guess it is a shaping error (e.g. haiku adds an extra dimension on Linear I think, so maybe that's messing it up?). Does the term work fine in isolation? (see my comment in #446 for debugging tips)