Fix bugs in Solver hash and eq.

Solver is used a static argument. Solver is polymorphic and the subclass identity is used as an implicit argument affecting the trace, so the subclass identity must be part of the key for the Jax cache. The Solver __hash__ and __eq__ methods are not hashing and comparing the class id. This CL fixes that by comparing a string representation of the class id, including the module path.

PiperOrigin-RevId: 819912183

Author: goodfeli

torax/_src/solver/solver.py

@@ -33,7 +33,7 @@
 
 
 class Solver(abc.ABC):
-  """Solves for a single time steps update to State.
+  """Solves for a single time step's update to State.
 
   Attributes:
     physics_models: Physics models.
@@ -45,11 +45,18 @@ def __init__(
   ):
     self.physics_models = physics_models
 
+  def _class_id(self) -> str:
+    return f'{self.__class__.__module__}.{self.__class__.__qualname__}'
+
   def __hash__(self) -> int:
-    return hash(self.physics_models)
+    return hash((self._class_id(), self.physics_models))
 
   def __eq__(self, other: typing_extensions.Self) -> bool:
-    return self.physics_models == other.physics_models
+    return (
+        isinstance(other, Solver)
+        and self._class_id() == other._class_id()
+        and self.physics_models == other.physics_models
+    )
 
   @functools.partial(
       jax.jit,