GitLab

* Adding JAX/Praxis modules and dependencies.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding UTs to JAX/Praxis modules.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Remove praxis as a dependency due to not strictly needed
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Repalce is_fp8_supported to is_fp8_available
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make Praxis as an optional dependency.

1. Removed 'from . import praxis' in __init__.py.
    1.1 Noted, keep 'from . import flax' for deprecated warning.
2. Changed te.flax to te_flax in examples and README.rst.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding a workaround to FP8 training on Praxis.
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adjust Module Structure.

1. Collect Flax related modules to a sub-folder, flax.
2. Add a function to unify scale_init for zero-centered-gamma LN.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Make changes be compatible to previous versions.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adapt jax/examples to the new module structure.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update jax/docs and Add deprecated warning.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Update README
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding deprecated_wrapper
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Adding deprecated warning to flax modules which imported via transformer_engine.jax
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Fix CI errors and update docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Removing unnecessary deprecated warning in docs.
Signed-off-by: Ming Huang <mingh@nvidia.com>

* Implementing __iter__ to DeprecatedEnum.
Signed-off-by: Ming Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming Huang <mingh@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* 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>

Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Add the link to the examples and the development user guide.
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>

* 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>

* Add tensorflow build

Improve build instructions

Fix pybind enum usage

Fix Python_EXECUTABLE cmake var

Move scale_inv calculations to FW
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Apply clang-format
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Format python files
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add TF build CI
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Lint checks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Another round of lint checks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix TF image tag
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>

* Use the existing recipe file
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add license claim blocks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix a bug about bias dtype conversion
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add mnist example and cleanup old examples
Signed-off-by: kaixih <kaixih@nvidia.com>

* Autopep8 the tests
Signed-off-by: kaixih <kaixih@nvidia.com>

* Autopep8 the examples
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add example in Readme
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add unit tests and linting for TensorFlow
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add causal mask for non-fused case
Signed-off-by: kaixih <kaixih@nvidia.com>

* Fix the mismatched TF vs TE masks
Signed-off-by: kaixih <kaixih@nvidia.com>

* Addressing CI tests
Signed-off-by: kaixih <kaixih@nvidia.com>

* Run lint test
Signed-off-by: kaixih <kaixih@nvidia.com>

* Add missing import
Signed-off-by: kaixih <kaixih@nvidia.com>

* Skip fp8 tests for pre-Hopper GPUs
Signed-off-by: kaixih <kaixih@nvidia.com>

* Remove non-pytest tests
Signed-off-by: kaixih <kaixih@nvidia.com>

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

---------
Signed-off-by: Trevor Morris <tmorris@nvidia.com>
Signed-off-by: kaixih <kaixih@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: kaixih <kaixih@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

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

* Adding JAX to README.rst
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Refine README.rst as the suggestion from review.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

* Refine the API doc of extend_logical_axis_rules.
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>

---------
Signed-off-by: Ming-Xu Huang <mingh@nvidia.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>

Fix the sample code so it compiles after the signature of `te.Linear` has changed.
Signed-off-by: nzmora-nvidia <96238833+nzmora-nvidia@users.noreply.github.com>
Signed-off-by: nzmora-nvidia <96238833+nzmora-nvidia@users.noreply.github.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>