Add nnx flax update (#114)

Author: jshin1394

src/metrax/nnx/nnx_wrapper.py

@@ -32,3 +32,6 @@ def update(self, **kwargs) -> None:
 
   def compute(self):
     return self.clu_metric.compute()
+
+  def __init_subclass__(cls, **kwargs):
+    super().__init_subclass__(pytree=False, **kwargs)