Fix `NVTE_FRAMEWORK=all` installation (#1850)

* Fix NVTE_FRAMEWORK=all
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Workflow tests and fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix jax install
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Update dep
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add numpy
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add dep
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/csrc/extensions/multi_tensor/l2norm.cpp

@@ -4,7 +4,7 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include "extensions.h"
+#include "../../extensions.h"
 
 namespace transformer_engine::pytorch {
 

transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp

@@ -4,7 +4,7 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include "extensions.h"
+#include "../../extensions.h"
 
 namespace transformer_engine::pytorch {
 

transformer_engine/pytorch/csrc/extensions/multi_tensor/sgd.cpp

@@ -4,7 +4,7 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include "extensions.h"
+#include "../../extensions.h"
 
 namespace transformer_engine::pytorch {
 

transformer_engine/pytorch/csrc/extensions/normalization.cpp

@@ -4,8 +4,8 @@
  * See LICENSE for license information.
  ************************************************************************/
 
+#include "../extensions.h"
 #include "common/util/system.h"
-#include "extensions.h"
 #include "pybind.h"
 
 namespace transformer_engine::pytorch {

transformer_engine/pytorch/csrc/extensions/padding.cpp

@@ -4,7 +4,7 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include "extensions.h"
+#include "../extensions.h"
 #include "pybind.h"
 
 namespace transformer_engine::pytorch {

transformer_engine/pytorch/csrc/extensions/permutation.cpp

@@ -4,7 +4,7 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include "extensions.h"
+#include "../extensions.h"
 
 namespace transformer_engine::pytorch {
 

transformer_engine/pytorch/csrc/extensions/recipe.cpp

@@ -9,7 +9,7 @@
 
 #include <string>
 
-#include "extensions.h"
+#include "../extensions.h"
 #include "transformer_engine/transformer_engine.h"
 
 namespace transformer_engine::pytorch {

transformer_engine/pytorch/csrc/extensions/softmax.cpp

@@ -4,7 +4,7 @@
  * See LICENSE for license information.
  ************************************************************************/
 
-#include "extensions.h"
+#include "../extensions.h"
 
 namespace transformer_engine::pytorch {
 

transformer_engine/pytorch/csrc/extensions/transpose.cpp

@@ -6,7 +6,7 @@
 
 #include <optional>
 
-#include "extensions.h"
+#include "../extensions.h"
 #include "pybind.h"
 
 namespace transformer_engine::pytorch {