From 2095ab861435085c61bf691ab123e37b3fd71781 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Dec 2025 02:40:22 +0000 Subject: [PATCH 1/2] Initial plan From ca56ef142fcc0ac59a6fe43362349323b2c2d5ca Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Dec 2025 02:42:57 +0000 Subject: [PATCH 2/2] Replace _graphdef with _ for unused variables Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com> --- tutorials/v1/device/08_initialize_nnx_on_device.py | 4 ++-- tutorials/v1/device/09_split_learning_vertical.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tutorials/v1/device/08_initialize_nnx_on_device.py b/tutorials/v1/device/08_initialize_nnx_on_device.py index 3a4b76aa..8fdd26fc 100644 --- a/tutorials/v1/device/08_initialize_nnx_on_device.py +++ b/tutorials/v1/device/08_initialize_nnx_on_device.py @@ -84,7 +84,7 @@ def init_model_logic(input_dim: int, hidden_dim: int, output_dim: int, seed: int ) # Split and convert state to pure dict - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) state_dict = state.to_pure_dict() print(f"[Device P0] Initialized model with {len(state_dict)} parameter groups") @@ -180,7 +180,7 @@ def init_model_with_optimizer_logic( ) # Split and convert model state to pure dict - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) model_state_dict = state.to_pure_dict() # Initialize optimizer and convert its state to pure dict diff --git a/tutorials/v1/device/09_split_learning_vertical.py b/tutorials/v1/device/09_split_learning_vertical.py index 68d141ea..e41252c8 100644 --- a/tutorials/v1/device/09_split_learning_vertical.py +++ b/tutorials/v1/device/09_split_learning_vertical.py @@ -646,14 +646,14 @@ def initialize_and_train_split_learning( def _init_alice_base(): model = AliceBaseModel(input_dim=m1, hidden_dim=h1, rngs=nnx.Rngs(seed_alice)) - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) return model_state_to_dict(state, opt_state, 0) def _init_bob_base(): model = BobBaseModel(input_dim=m2, hidden_dim=h2, rngs=nnx.Rngs(seed_bob)) - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) return model_state_to_dict(state, opt_state, 0) @@ -665,7 +665,7 @@ def _init_alice_agg(): output_dim=n_classes, rngs=nnx.Rngs(seed_agg), ) - _graphdef, state = nnx.split(model) + _, state = nnx.split(model) tx = optax.sgd(learning_rate) opt_state = tx.init(state.to_pure_dict()) return model_state_to_dict(state, opt_state, 0)