Comment by rlupi

Have you tried Jax or Taichi? https://www.taichi-lang.org/

For Monte Carlo simulations, Pyro and tensorflow_probability have also nice abstractions.