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

PiperOrigin-RevId: 700716670
google · Nov 27, 2024 · 5d896bc · 5d896bc
2 parents be2f861 + 7b50ffe
commit 5d896bc
Show file tree

Hide file tree

Showing 6 changed files with 146 additions and 206 deletions.
diff --git a/flax/nnx/reprlib.py b/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