flax

761a27d3 - Add tree-mode support to nnx.{cond,switch,eval_shape,checkify}

Add tree-mode support to nnx.{cond,switch,eval_shape,checkify} Adds `graph` parameter to `nnx.cond`, `nnx.switch`, `nnx.eval_shape` and `nnx.checkify` allowing them to operate in tree-mode when `graph=False`. Existing tests are parametrized to run with both graph and tree mode. PiperOrigin-RevId: 873032717