Is there a reason you are limited to JAX/tensorflow and can't use pytorch?

For a lot of people I know whose main job was not to write code, switching from tensorflow to pytorch was something that saved them ten to hundreds of hours in the long run, even accounting for the initial learning time.