Part 3: Transformations# This section covers JAX transformations and program optimization in BrainState. JIT Compilation Automatic Differentiation Vectorization Parallelisation Loops and Conditionals Other Transforms
