чи є у когось досвід навчання (у масштабі) як у Jax, так і в Pytorch? якщо так, то в якому масштабі ви дійсно знаходите больові удари в Pytorch і в якому масштабі, на вашу думку, загальний біль Jax однозначно того вартий?