GitLab

* Clean up the installation instruction. We where telling to install the dev version in the README.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Typos
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

* Update installation instructio for JAX and add some depenencies.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Bring back support for none pip installed pybind11.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Apply suggestions from code review
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>

* Changes following review.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Change order to make it more clear.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Add other reviers suggestion.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* pybind11 is needed for all FW.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Add flax as a dep
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Update README.rst
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>

---------
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Signed-off-by: Frédéric Bastien <frederic.bastien@gmail.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>