Internal change

PiperOrigin-RevId: 408236164

tensorflow_models/__init__.py

@@ -16,3 +16,8 @@
 # pylint: disable=wildcard-import
 from tensorflow_models import nlp
 from tensorflow_models import vision
+
+from official import core
+from official.modeling import hyperparams
+from official.modeling import optimization
+from official.modeling import tf_utils

tensorflow_models/tensorflow_models_test.py

@@ -20,13 +20,20 @@ import tensorflow_models as tfm
 
 class TensorflowModelsTest(tf.test.TestCase):
 
-  def testImport(self):
-    _ = tfm.nlp.layers.TransformerEncoderBlock(
-        num_attention_heads=2, inner_dim=10, inner_activation='relu')
+  def testVisionImport(self):
     _ = tfm.vision.layers.SqueezeExcitation(
         in_filters=8, out_filters=4, se_ratio=1)
     _ = tfm.vision.configs.image_classification.Losses()
 
+  def testNLPImport(self):
+    _ = tfm.nlp.layers.TransformerEncoderBlock(
+        num_attention_heads=2, inner_dim=10, inner_activation='relu')
+
+  def testCommonImports(self):
+    _ = tfm.hyperparams.Config()
+    _ = tfm.optimization.LinearWarmup(
+        after_warmup_lr_sched=0.0, warmup_steps=10, warmup_learning_rate=0.1)
+
 
 if __name__ == '__main__':
   tf.test.main()