Add MetaEnum to get desired behavior of in-operator

curvlinops/kfac.py

@@ -19,7 +19,7 @@
 from __future__ import annotations
 
 from collections.abc import MutableMapping
-from enum import Enum
+from enum import Enum, EnumMeta
 from functools import partial
 from math import sqrt
 from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, TypeVar, Union
@@ -52,7 +52,18 @@
 ParameterMatrixType = TypeVar("ParameterMatrixType", Tensor, List[Tensor])
 
 
-class FisherType(str, Enum):
+class MetaEnum(EnumMeta):
+    """Metaclass for the Enum class for desired behavior of the `in` operator."""
+
+    def __contains__(cls, item):
+        try:
+            cls(item)
+        except ValueError:
+            return False
+        return True
+
+
+class FisherType(str, Enum, metaclass=MetaEnum):
     """Enum for the Fisher type."""
 
     TYPE2 = "type-2"
@@ -61,7 +72,7 @@ class FisherType(str, Enum):
     FORWARD_ONLY = "forward-only"
 
 
-class KFACType(str, Enum):
+class KFACType(str, Enum, metaclass=MetaEnum):
     """Enum for the KFAC approximation type."""
 
     EXPAND = "expand"