Firstly vmap() might make it easier for you from a programming standpoint. You can think of operations in terms of single instances and then autobatch them. Al

Firstly vmap() might make it easier for you from a programming standpoint. You can think of operations in terms of single instances and then autobatch them. Also, jit() is where you'll likely see your speed increases. As for autodiff - not every application will use it. What exactly are you trying to optimize? If you can parameterize your model somehow (e.g. parameterize edge weights) and then figure out some way to measure how "bad" your model is, you can use autodiff to tune your edge weights to minimize that metric. Not too familiar with DAGs but the first step is figuring out how to parameterize your model in a way that tuning the params can lead to your goal (and how to measure how close you are to that goal)