From 7f970677e40730a6e25066da684167ffccf99e43 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Fri, 17 Oct 2025 23:49:04 +0200 Subject: [PATCH 1/2] Added Enumeration.string --- equinox/_enum.py | 50 ++++++++++++++++++++++++++++++++++++++++------ tests/test_enum.py | 16 ++++++++++++++- 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/equinox/_enum.py b/equinox/_enum.py index 70a3fc3d..fd5f22af 100644 --- a/equinox/_enum.py +++ b/equinox/_enum.py @@ -59,9 +59,15 @@ def __new__(mcs, cls_name, bases, namespace): _base_offsets=base_offsets, ) if cls_name == "Enumeration" and "Enumeration" not in globals().keys(): + new_namespace["string"] = namespace.pop("string") new_namespace["promote"] = namespace.pop("promote") new_namespace["where"] = namespace.pop("where") else: + if "string" in namespace: + raise ValueError( + "Cannot have enumeration item with name `string`, as this " + "conflicts with the classmethod of this name" + ) if "promote" in namespace: raise ValueError( "Cannot have enumeration item with name `promote`, as this " @@ -207,16 +213,17 @@ def is_traced(self) -> bool: from typing import ClassVar from typing_extensions import Self - class _Sequence(type): - def __getitem__(cls, item) -> str: ... - + class _HasLen(type): def __len__(cls) -> int: ... - class Enumeration(enum.Enum, EnumerationItem, metaclass=_Sequence): # pyright: ignore + class Enumeration(enum.Enum, EnumerationItem, metaclass=_HasLen): # pyright: ignore _name_to_item: ClassVar[dict[str, EnumerationItem]] # pyright: ignore _index_to_message: ClassVar[list[str]] # pyright: ignore _base_offsets: ClassVar[dict["Enumeration", int]] + @classmethod + def string(cls, item: "Enumeration") -> str: ... + @classmethod def promote(cls, item: "Enumeration") -> Self: ... @@ -262,10 +269,10 @@ class RESULTS(eqx.Enumeration): print(RESULTS.success) # RESULTS ``` - Given a Enumeration element, just the string can be looked up by indexing it: + Given an Enumeration element, just the string message can be looked up: ```python result = RESULTS.success - print(RESULTS[result]) # Hurrah! + print(RESULTS.string(result)) # Hurrah! ``` Enumerations support inheritance, to include all of the superclasses' fields as @@ -292,6 +299,37 @@ class MORE_RESULTS(RESULTS): # pyright: ignore string corresponding to the enumeration item. """ + @classmethod + def string(cls, item: _Enumeration) -> str: + """Look up the string associated with an enum element. + + !!! Example + + ```python + class RESULTS(eqx.Enumeration): + success = "success" + linear_solve_failed = "Linear solve failed to converge" + + result = RESULTS.linear_solve_failed + assert RESULTS.string(result) == "Linear solve failed to converge" + ``` + + **Arguments:** + + - `item`: a member of the enumeration. + + **Returns:** + + If `item` is outside of JIT, then the string associated with the enum (from + the right hand side of the equals sign during its definition). + + If `item` is a traced JAX value (inside of JIT), then the string "traced" is + returned. + """ + # Wrap old API with new API, see + # https://github.com/patrick-kidger/equinox/issues/1121. + return cls[item] + @classmethod def promote(cls, item: _Enumeration) -> _Enumeration: """Enums support `.promote` (on the class) to promote from an inherited diff --git a/tests/test_enum.py b/tests/test_enum.py index 6e2d0cb8..a06806b9 100644 --- a/tests/test_enum.py +++ b/tests/test_enum.py @@ -207,13 +207,27 @@ class A(eqxi.Enumeration): class B(eqxi.Enumeration): b = "bye" - A[A.a] + assert A[A.a] == "hi" with pytest.raises(ValueError): A[0] with pytest.raises(ValueError): A[B.b] +def test_string(): + class A(eqxi.Enumeration): + a = "hi" + + class B(eqxi.Enumeration): + b = "bye" + + assert A.string(A.a) == "hi" + with pytest.raises(ValueError): + A.string(0) # pyright: ignore[reportArgumentType] + with pytest.raises(ValueError): + A.string(B.b) + + def test_isinstance(): class A(eqxi.Enumeration): a = "hi" From bf9f3a2e3824718b21d5eefef51a808c4b5ea474 Mon Sep 17 00:00:00 2001 From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com> Date: Fri, 17 Oct 2025 23:51:32 +0200 Subject: [PATCH 2/2] Pre-commit fixes for JAX 0.8.0 --- equinox/_filters.py | 4 +++- equinox/_serialisation.py | 2 +- tests/test_serialisation.py | 8 ++++---- tests/test_tree.py | 6 +++--- 4 files changed, 11 insertions(+), 9 deletions(-) diff --git a/equinox/_filters.py b/equinox/_filters.py index 624a976f..51114f71 100644 --- a/equinox/_filters.py +++ b/equinox/_filters.py @@ -22,7 +22,9 @@ _array_types = (np.ndarray, np.generic, jax.Array) # JAX < 0.7.2 try: # JAX 0.7.2 - from jax._src.literals import LiteralArray + from jax._src.literals import ( + LiteralArray, # pyright: ignore[reportAttributeAccessIssue] + ) _array_types += (LiteralArray,) except ImportError: diff --git a/equinox/_serialisation.py b/equinox/_serialisation.py index 64ab9077..56dcb10f 100644 --- a/equinox/_serialisation.py +++ b/equinox/_serialisation.py @@ -135,7 +135,7 @@ def default_deserialise_filter_spec(f: BinaryIO, x: Any) -> Any: out = np.load(f) if isinstance(x, jax.dtypes.bfloat16): out = out.view(jax.dtypes.bfloat16) - return type(x)(out.item()) + return type(x)(out.item()) # pyright: ignore[reportCallIssue] else: return x diff --git a/tests/test_serialisation.py b/tests/test_serialisation.py index 44b214ba..a1ac523f 100644 --- a/tests/test_serialisation.py +++ b/tests/test_serialisation.py @@ -122,22 +122,22 @@ def test_helpful_errors(getkey, tmp_path): def test_generic_dtype_serialisation(getkey, tmp_path): # Ensure we can round trip when we start with an array - jax_array = jnp.array(bfloat16(1)) + jax_array = jnp.array(bfloat16(1)) # pyright: ignore[reportCallIssue] eqx.tree_serialise_leaves(tmp_path, jax_array) - like_jax_array = jnp.array(bfloat16(2)) + like_jax_array = jnp.array(bfloat16(2)) # pyright: ignore[reportCallIssue] loaded_jax_array = eqx.tree_deserialise_leaves(tmp_path, like_jax_array) assert jax_array.item() == loaded_jax_array.item() tree = ( jnp.array(1e-8), - bfloat16(1e-8), + bfloat16(1e-8), # pyright: ignore[reportCallIssue] np.float32(1e-8), jnp.array(1e-8), np.float64(1e-8), ) like_tree = ( jnp.array(2.0), - bfloat16(2), + bfloat16(2), # pyright: ignore[reportCallIssue] np.float32(2), jnp.array(2.0), np.float64(2.0), diff --git a/tests/test_tree.py b/tests/test_tree.py index b4cae5b1..564f5057 100644 --- a/tests/test_tree.py +++ b/tests/test_tree.py @@ -238,9 +238,9 @@ def test_tree_equal_scalars(): # dtype does matter assert _typeequal(eqx.tree_equal(x, z), False) - z = jax.dtypes.bfloat16(1) - z2 = jax.dtypes.bfloat16(1) - w = jax.dtypes.bfloat16(2) + z = jax.dtypes.bfloat16(1) # pyright: ignore[reportCallIssue] + z2 = jax.dtypes.bfloat16(1) # pyright: ignore[reportCallIssue] + w = jax.dtypes.bfloat16(2) # pyright: ignore[reportCallIssue] assert _typeequal(eqx.tree_equal(z, z2), True) assert _typeequal(eqx.tree_equal(z, w), False)