Skip to content

Address review feedback: clarify documentation and fix type consistency - #306

Merged
da-niao-dan merged 3 commits into
peicheng/nnx_split_learning_tutorialfrom
copilot/sub-pr-304
Dec 19, 2025
Merged

da-niao-dan merged 3 commits into
peicheng/nnx_split_learning_tutorialfrom
copilot/sub-pr-304

Conversation

Copilot AI commented Dec 19, 2025

Copy link
Copy Markdown
Contributor

Addresses review comments on tutorials 08 and 09 by correcting misleading documentation and fixing a type inconsistency bug.

Tutorial 08: Fix nnx.eval_shape documentation

  • Removed inaccurate "memory efficient" claims - nnx.eval_shape creates abstract structures but parameters are materialized immediately by replace_by_pure_dict
  • Updated comments at lines 121, 237, and key takeaways to reflect actual behavior

Tutorial 09: Multiple improvements

Removed non-existent pattern from documentation:

  • Deleted nnx.eval_shape from key concepts and takeaways - pattern not demonstrated in tutorial code

Clarified reconstruction requirements:

  • Documented that graphdef reconstruction requires exact same initialization parameters (input_dim, hidden_dim, output_dim, seed)
  • Explained parameter capture vs module-level constants in JAX tracing context

Fixed type bug:

# Before: type mismatch - Python float if n_steps=0, JAX array if n_steps>0
final_loss = None  # or 0.0

# After: consistent JAX array type
final_loss = jnp.array(0.0)

Improved function naming:

  • train_split_learning_for_n_stepsinitialize_and_train_split_learning (performs both init and training)

Enhanced dtype workaround documentation:

  • Added detailed explanation of table_to_tensor dtype constraint with alternative patterns for non-integer labels

💡 You can make Copilot smarter by setting up custom instructions, customizing its development environment and configuring Model Context Protocol (MCP) servers. Learn more Copilot coding agent tips in the docs.

Copilot AI and others added 2 commits December 19, 2025 02:24
Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com>
Co-authored-by: da-niao-dan <9532472+da-niao-dan@users.noreply.github.com>
Copilot AI changed the title [WIP] Add more NNX and SL tutorials for MPLang v1 Address review feedback: clarify documentation and fix type consistency Dec 19, 2025
Copilot AI requested a review from da-niao-dan December 19, 2025 02:28

@da-niao-dan da-niao-dan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@da-niao-dan
da-niao-dan marked this pull request as ready for review December 19, 2025 02:33
@da-niao-dan
da-niao-dan requested a review from a team as a code owner December 19, 2025 02:33
@da-niao-dan
da-niao-dan merged commit ad3afbc into peicheng/nnx_split_learning_tutorial Dec 19, 2025
6 checks passed
@da-niao-dan
da-niao-dan deleted the copilot/sub-pr-304 branch December 19, 2025 02:33
@github-actions github-actions Bot locked and limited conversation to collaborators Dec 19, 2025
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants