Merge pull request #332 from ROCm/genesu/fix-small-bugs

[Fix] Missed import and dim reference

Author: GeneDer

setup.py

@@ -152,6 +152,8 @@ def setup_requirements() -> Tuple[List[str], List[str], List[str]]:
     if not found_cmake():
         setup_reqs.append("cmake>=3.21")
     if not found_ninja():
+        import sys
+
         subprocess.check_call([sys.executable, "-m", "pip", "install", "ninja"])
         setup_reqs.append("ninja")
     if not found_pybind11():

transformer_engine/jax/flax/transformer.py

@@ -637,7 +637,7 @@ def __call__(
             dropout_rng = self.make_rng(self.dropout_rng_name)
 
         if self.scale_factor is None:
-            scale_factor = 1.0 / sqrt(self.head_dim_qk)
+            scale_factor = 1.0 / sqrt(head_dim_qk)
         else:
             scale_factor = self.scale_factor
         del self.scale_factor