Fix import statements in README#15

RaulPL wants to merge 1 commit into

probml:mainprobml/sts-jax:mainfrom

RaulPL:fix_components_in_readmeRaulPL/sts-jax:fix_components_in_readmeCopy head branch name to clipboard