Update installation docs (#472)

Author: mthrok

docs/source/installation.rst

@@ -61,6 +61,7 @@ Requirements
     - NumPy
     - (Optional) PyTorch
     - (Optional) Numba
+    - (Optional) Jax
     - (Optional) CUDA runtime
 
 Dependencies