[fix] Add check for newer versions of jax#3350

alberttorosyan wants to merge 2 commits intomainaimhubio/aim:mainfrom

fix/jax-0.6.2-arrays-supportaimhubio/aim:fix/jax-0.6.2-arrays-supportCopy head branch name to clipboard