Merge pull request #4400 from google:nnx-more-optimizations-p2

PiperOrigin-RevId: 700716670

Author: Flax Authors

flax/nnx/reprlib.py

@@ -16,7 +16,6 @@
 import dataclasses
 import threading
 import typing as tp
-from abc import ABC, abstractmethod
 
 A = tp.TypeVar('A')
 B = tp.TypeVar('B')
@@ -48,10 +47,9 @@ class Attr:
   end: str = ''
 
 
-class Representable(ABC):
+class Representable:
   __slots__ = ()
 
-  @abstractmethod
   def __nnx_repr__(self) -> tp.Iterator[tp.Union[Object, Attr]]:
     raise NotImplementedError