Merge branch 'main' of https://github.com/NVIDIA/TransformerEngine

transformer_engine/jax/csrc/extensions/misc.h

@@ -34,5 +34,11 @@ inline size_t product(const std::vector<size_t> &shape) {
   return ret;
 }
 
+enum class QuantizeAxis {
+  ROWWISE,
+  COLWISE,
+  ROWWISE_COLWISE,
+};
+
 }  // namespace jax
 }  // namespace transformer_engine

transformer_engine/jax/flax/__init__.py

@@ -3,7 +3,7 @@
 # See LICENSE for license information.
 """Transformer Engine bindings for JAX"""
 from .module import DenseGeneral, LayerNorm
-from .module import LayerNormDenseGeneral, LayerNormMLP, TransformerEngineBase
+from .module import LayerNormDenseGeneral, LayerNormMLP
 from .transformer import extend_logical_axis_rules
 from .transformer import DotProductAttention, MultiHeadAttention, RelativePositionBiases
 from .transformer import TransformerLayer, TransformerLayerType
@@ -13,7 +13,6 @@ __all__ = [
     "LayerNorm",
     "LayerNormDenseGeneral",
     "LayerNormMLP",
-    "TransformerEngineBase",
     "extend_logical_axis_rules",
     "DotProductAttention",
     "MultiHeadAttention",