可视化

flax.nnx.display(*args)

使用 Treescope 优美格式化打印工具显示给定的对象。

如果未安装 treescope 或代码未在 IPython 中运行，display 将改为打印对象。