Add bert_padding

__init__.py

@@ -34,7 +34,10 @@ from .model import (
     FlexBertForSequenceClassification,
     FlexBertForMultipleChoice,
 )
-
+from .bert_padding import(
+    IndexFirstAxis,
+    IndexPutFirstAxis
+)
 
 __all__ = [
     "BertAlibiEmbeddings",
@@ -65,4 +68,6 @@ __all__ = [
     "FlexBertForMaskedLM",
     "FlexBertForSequenceClassification",
     "FlexBertForMultipleChoice",
+    "IndexFirstAxis,
+    "IndexPutFirstAxis"
 ]