I have been working on my DNN model using TensorFlow even though ML is not my main research. But it is a substantial part of my research, so I have to figure things out on my own, and I have done so over the past 3 years. However, I spend so much time on figuring out how any of TF methods works and debugging them. I never used JAX but I am not sure if this sort of grinding is normal when you use JAX as well (I always hear great things about JAX). I have built so many things using TF I don't think it is wise for me to learn JAX and migrate my work into JAX code base.