AutoOrderedAcrossTypes

AutoOrderedAcrossTypes#

class penzai.core.auto_order_types.AutoOrderedAcrossTypes[source]#

Bases: object

Mixin to define an arbitrary total ordering across dataclass subclasses.

Subclasses of this class will be comparable with other subclasses of this class, using an arbitrary total ordering of their types. Additionally, any subclass of this class will be greater than any string or tuple. This makes it possible to use subclasses of this class as keys in dictionaries that are passed through JAX transformations, potentially in combination with string or tuple keys.

Every subclass of this class must be a dataclass with eq=True, order=False (the default), otherwise the provided comparisons may not work correctly. (Setting eq=False may cause inconsistencies between equality and ordering comparisons, and setting order=True will override the orderings defined here.)

An example of how to subclass this class for the common case of making a new dictionary key type:

@dataclasses.dataclass(frozen=True)  # WITHOUT order=True
class MyType(auto_order_types.AutoOrderedAcrossTypes):
  some_field: int
  some_other_field: str

Here MyType will have the same ordering behavior as it would have if order=True was set, but it will also be comparable with other subclasses of AutoOrderedAcrossTypes.

Inherited Methods

(expand to view inherited methods)

__init__()

Contents