added @DataClass so that lint check does not complain about field()

Signed-off-by: Phuong Nguyen <[email protected]>

transformer_engine/jax/praxis/module.py

@@ -6,7 +6,7 @@
 """
 from functools import partial
 from typing import Callable, Iterable, Sequence, Tuple, Union
-from dataclasses import field
+from dataclasses import field, dataclass
 
 from praxis import pax_fiddle
 from praxis.base_layer import init_var
@@ -28,6 +28,7 @@ def _generate_ln_scale_init(scale_init):
     return scale_init
 
 
+@dataclass
 class TransformerEngineBaseLayer(BaseLayer):
     """TransformerEngineBaseLayer"""
 
@@ -67,6 +68,7 @@ def create_layer(self, name, flax_module_cls):
         self.create_child(name, flax_module_p.clone())
 
 
+@dataclass
 class LayerNorm(TransformerEngineBaseLayer):
     """LayerNorm"""
 
@@ -103,6 +105,7 @@ def __call__(self, x: JTensor) -> JTensor:
         return self.layer_norm(x)
 
 
+@dataclass
 class FusedSoftmax(TransformerEngineBaseLayer):
     """FusedSoftmax"""
 
@@ -124,6 +127,7 @@ def __call__(self, x: JTensor, mask: JTensor = None, bias: JTensor = None) -> JT
         return self.fused_softmax(x, mask, bias)
 
 
+@dataclass
 class Linear(TransformerEngineBaseLayer):
     """Linear"""
 
@@ -165,6 +169,7 @@ def __call__(self, x: JTensor) -> JTensor:
         return self.linear(x)
 
 
+@dataclass
 class LayerNormLinear(TransformerEngineBaseLayer):
     """LayerNormLinear"""
 
@@ -228,6 +233,7 @@ def __call__(self, x: JTensor) -> JTensor:
         return self.ln_linear(x)
 
 
+@dataclass
 class LayerNormMLP(TransformerEngineBaseLayer):
     """LayerNormMLP"""
 

transformer_engine/jax/praxis/transformer.py

@@ -6,7 +6,7 @@
 """
 from functools import partial
 from typing import Optional, Sequence, Tuple
-from dataclasses import field
+from dataclasses import field, dataclass
 import warnings
 
 from praxis import pax_fiddle
@@ -22,6 +22,7 @@
 from ..attention import AttnBiasType, AttnMaskType
 
 
+@dataclass
 class RelativePositionBiases(TransformerEngineBaseLayer):
     """RelativePositionBiases"""
 
@@ -67,6 +68,7 @@ def __call__(self, q_seqlen: JTensor, k_seqlen: JTensor, bidirectional: bool = T
         return self.relative_position_bias(q_seqlen, k_seqlen, bidirectional)
 
 
+@dataclass
 class DotProductAttention(TransformerEngineBaseLayer):
     """DotProductAttention"""
 
@@ -125,6 +127,7 @@ def __call__(
         )
 
 
+@dataclass
 class MultiHeadAttention(TransformerEngineBaseLayer):
     """MultiHeadAttention"""
 
@@ -258,6 +261,7 @@ def __call__(
         )
 
 
+@dataclass
 class TransformerLayer(TransformerEngineBaseLayer):
     """TransformerLayer"""