from flax import nnx

class P(nnx.Variable):
  def on_set_value(self, value):
    return self.value + 1.0

class M(nnx.Module):
  def __init__(self):
    self.param = P(value=0.0)

  def __call__(self, x):
    return self.param * x

nnx.tabulate(M(), 2.0)
# ---------------------------------------------------------------------------
# RepresenterError                          Traceback (most recent call last)
# [/tmp/ipython-input-2062670779.py](https://qy6gqpckbx-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20251003-060040_RC00_814643377#) in <cell line: 0>()
#      12     return self.param * x
#      13 
# ---> 14 nnx.tabulate(M(), 2.0)
# 
# 12 frames
# [/usr/local/lib/python3.12/dist-packages/yaml/representer.py](https://qy6gqpckbx-496ff2e9c6d22116-0-colab.googleusercontent.com/outputframe.html?vrz=colab_20251003-060040_RC00_814643377#) in represent_undefined(self, data)
#     229 
#     230     def represent_undefined(self, data):
# --> 231         raise RepresenterError("cannot represent an object", data)
#     232 
#     233 SafeRepresenter.add_representer(type(None),
#
# RepresenterError: ('cannot represent an object', <function P.on_set_value at 0x789fe827d3a0>)