Adds a helper function.

Author: mark

axlearn/common/config.py

@@ -64,6 +64,7 @@ class Config(ConfigBase):
 import copy
 import dataclasses
 import enum
+import importlib
 import inspect
 import re
 import types
@@ -255,28 +256,28 @@ def validate_config_field_name(name: str) -> None:
     match_fn=lambda v: not isinstance(v, type) and hasattr(v, "from_pretrained"),
     validate_fn=lambda v: validate_config_field_value(v.to_dict()),
 )
-# Register other validators for convenience and backwards compat.
-# We allow these to be optional to avoid a hard dependency.
-try:
-    import numpy as np
 
-    register_validator(
-        match_fn=lambda v: isinstance(v, np.dtype),
-        validate_fn=lambda _: None,
-    )
-except ImportError:
-    pass
 
-try:
-    # As of 0.6.1, PartitionSpec is not longer a tuple.
-    import jax
+def _maybe_register_optional_type(module: str, attribute: str):
+    """Attempts to register a valid type specified via `module` path and `attribute` name.
 
-    register_validator(
-        match_fn=lambda v: isinstance(v, jax.sharding.PartitionSpec),
-        validate_fn=lambda _: None,
-    )
-except ImportError:
-    pass
+    This is used to register optional types so that we can avoid a hard dependency on the module.
+    """
+    try:
+        module = importlib.import_module(module)
+        register_validator(
+            match_fn=lambda v: isinstance(v, getattr(module, attribute)),
+            validate_fn=lambda _: None,
+        )
+    except (ImportError, AttributeError):
+        pass
+
+
+# Register other validators for convenience and backwards compat.
+# We allow these to be optional to avoid a hard dependency.
+_maybe_register_optional_type("numpy", "dtype")
+# As of 0.6.1, PartitionSpec is not longer a tuple.
+_maybe_register_optional_type("jax.sharding", "PartitionSpec")
 
 
 def validate_config_field_value(value: Any) -> None:

axlearn/common/config_test.py

@@ -1027,6 +1027,7 @@ class Config(Configurable.Config):
 
     def test_custom_validators(self):
         try:
+            # pytype: disable=invalid-annotation
             # pylint: disable-next=import-outside-toplevel
             import jax
 
@@ -1037,6 +1038,7 @@ class Test(ConfigBase):
             spec = jax.sharding.PartitionSpec("test")
             cfg = Test(partition_spec=spec)
             self.assertEqual(cfg.partition_spec, spec)
+            # pytype: enable=invalid-annotation
 
         except ImportError:
             pass