Update backend import in JAX speedup notebook#1

che-shr-cat merged 1 commit intoche-shr-cat:mainche-shr-cat/Deep_Learning_with_JAX:mainfrom

jimangel:patch-1jimangel/Deep_Learning_with_JAX:patch-1Copy head branch name to clipboard