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

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

alberttorosyan wants to merge 2 commits into

mainaimhubio/aim:mainfrom

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