docs: add JAX documentation link and update documentation dependencies

docs: add JAX documentation link and update documentation dependencies #94