Ensure correct module path when calling nnx.

[liudangyi]

Qwix relies on the linen module paths to apply the correct quantization configs. The module path is broken when part of the model is converted to nnx. This patch introduces a mechanism to update the linen module paths when calling a nnx module.

This patch also disables the quantization of dot-product attentions, which follows the same behavior as before.

For model_name=default, before:

[QWIX] module='decoder/layers/self_attention/query' op=dot_general0 rule=0
[QWIX] module='decoder/layers/self_attention/key' op=dot_general0 rule=0
[QWIX] module='decoder/layers/self_attention/value' op=dot_general0 rule=0
[QWIX] module='decoder/layers/self_attention/attention_op' op=einsum0 rule=0
[QWIX] module='decoder/layers/self_attention/attention_op' op=einsum1 rule=0
[QWIX] module='decoder/layers/self_attention/out' op=dot_general0 rule=0
[QWIX] module='decoder/layers/mlp' op=dot_general0 rule=0
[QWIX] module='decoder/layers/mlp' op=dot_general1 rule=0
[QWIX] module='decoder/layers/mlp' op=dot_general2 rule=0
[QWIX] module='decoder/logits_dense' op=dot_general0 rule=None

after:

[QWIX] module='decoder/layers/self_attention/query' op=dot_general0 rule=1
[QWIX] module='decoder/layers/self_attention/key' op=dot_general0 rule=1
[QWIX] module='decoder/layers/self_attention/value' op=dot_general0 rule=1
[QWIX] module='decoder/layers/self_attention/attention_op' op=einsum0 rule=0
[QWIX] module='decoder/layers/self_attention/attention_op' op=einsum1 rule=0
[QWIX] module='decoder/layers/self_attention/out' op=dot_general0 rule=1
[QWIX] module='decoder/layers/mlp/wi_0' op=dot_general0 rule=1
[QWIX] module='decoder/layers/mlp/wi_1' op=dot_general1 rule=1
[QWIX] module='decoder/layers/mlp/wo' op=dot_general2 rule=1
[QWIX] module='decoder/logits_dense' op=dot_general0 rule=None

For model_name=deepseek3-671b, before:

[QWIX] module='decoder/dense_layers/self_attention/wq_a' op=dot_general0 rule=0
[QWIX] module='decoder/dense_layers/self_attention/wq_b' op=dot_general0 rule=0
[QWIX] module='decoder/dense_layers/self_attention/wkv_a' op=dot_general0 rule=0
[QWIX] module='decoder/dense_layers/self_attention/wkv_b' op=dot_general0 rule=0
[QWIX] module='decoder/dense_layers/self_attention/attention_op' op=einsum0 rule=0
[QWIX] module='decoder/dense_layers/self_attention/attention_op' op=einsum1 rule=0
[QWIX] module='decoder/dense_layers/self_attention/out' op=dot_general0 rule=0
[QWIX] module='decoder/dense_layers/mlp' op=dot_general0 rule=0
[QWIX] module='decoder/dense_layers/mlp' op=dot_general1 rule=0
[QWIX] module='decoder/dense_layers/mlp' op=dot_general2 rule=0
[QWIX] module='decoder/moe_layers/self_attention/wq_a' op=dot_general0 rule=0
[QWIX] module='decoder/moe_layers/self_attention/wq_b' op=dot_general0 rule=0
[QWIX] module='decoder/moe_layers/self_attention/wkv_a' op=dot_general0 rule=0
[QWIX] module='decoder/moe_layers/self_attention/wkv_b' op=dot_general0 rule=0
[QWIX] module='decoder/moe_layers/self_attention/attention_op' op=einsum0 rule=0
[QWIX] module='decoder/moe_layers/self_attention/attention_op' op=einsum1 rule=0
[QWIX] module='decoder/moe_layers/self_attention/out' op=dot_general0 rule=0
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/MoeBlock_0/gate' op=dot_general0 rule=0
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/MoeBlock_0' op=einsum0 rule=0
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/shared_experts' op=dot_general0 rule=0
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/shared_experts' op=dot_general1 rule=0
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/shared_experts' op=dot_general2 rule=0
[QWIX] module='decoder/logits_dense' op=dot_general0 rule=None

after:

[QWIX] module='decoder/dense_layers/self_attention/wq_a' op=dot_general0 rule=1
[QWIX] module='decoder/dense_layers/self_attention/wq_b' op=dot_general0 rule=1
[QWIX] module='decoder/dense_layers/self_attention/wkv_a' op=dot_general0 rule=1
[QWIX] module='decoder/dense_layers/self_attention/wkv_b' op=dot_general0 rule=1
[QWIX] module='decoder/dense_layers/self_attention/attention_op' op=einsum0 rule=0
[QWIX] module='decoder/dense_layers/self_attention/attention_op' op=einsum1 rule=0
[QWIX] module='decoder/dense_layers/self_attention/out' op=dot_general0 rule=1
[QWIX] module='decoder/dense_layers/mlp/wi_0' op=dot_general0 rule=1
[QWIX] module='decoder/dense_layers/mlp/wi_1' op=dot_general1 rule=1
[QWIX] module='decoder/dense_layers/mlp/wo' op=dot_general2 rule=1
[QWIX] module='decoder/moe_layers/self_attention/wq_a' op=dot_general0 rule=1
[QWIX] module='decoder/moe_layers/self_attention/wq_b' op=dot_general0 rule=1
[QWIX] module='decoder/moe_layers/self_attention/wkv_a' op=dot_general0 rule=1
[QWIX] module='decoder/moe_layers/self_attention/wkv_b' op=dot_general0 rule=1
[QWIX] module='decoder/moe_layers/self_attention/attention_op' op=einsum0 rule=0
[QWIX] module='decoder/moe_layers/self_attention/attention_op' op=einsum1 rule=0
[QWIX] module='decoder/moe_layers/self_attention/out' op=dot_general0 rule=1
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/MoeBlock_0/gate' op=dot_general0 rule=1
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/MoeBlock_0' op=einsum0 rule=1
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/shared_experts/wi_0' op=dot_general0 rule=1
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/shared_experts/wi_1' op=dot_general1 rule=1
[QWIX] module='decoder/moe_layers/DeepSeekMoeBlock_0/shared_experts/wo' op=dot_general2 rule=1
[QWIX] module='decoder/logits_dense' op=dot_general0 rule=None

I also tested with PR #2066 (migrating attention to nnx) and the module paths are the same.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.