Improve error message for when nnx.Modules use jax or numpy arrays as leaf values

[RaghuSpaceRajan]

What does this PR do?

Fixes #4480

Checklist

• [✔] This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).

[cgarciae]

@RaghuSpaceRajan updated main today to fix the broken tests due to JAX's recent update. Can you please rebase?

[RaghuSpaceRajan]

@cgarciae Done.

[RaghuSpaceRajan]

Actually, for some reason, it shows me as author for some of the rebased commits. Not sure if that's normal, it's the first time I have rebased across forks. Let me know if should just redo the thing.

[cgarciae]

Hey @RaghuSpaceRajan , the diff looks good here on github.

[RaghuSpaceRajan]

Hi @cgarciae , sorry I was on vacation last week. Some conflict seems to have appeared in the meantime. I added your suggestion and tried to resolve the conflict but I am unsure about two things:

1. The join in the "path is not None" else part of the if statement here

Old code:

        path_str = '/'.join(map(str, (*path, key)))

New code:

          path_str = '/'.join(map(str, path))

Why is the key not needed anymore and why does path not have * anymore?

2. I hope error message in the part of the code that is in the new "path is None" else part of the if statement here makes sense even with path being None.

Are the changes good and should we merge now?

[cgarciae]

Hi @RaghuSpaceRajan , the reason is that path is constructed a bit differently and the key is appended early.

Thanks for updating the PR, I'll run the tests.

[RaghuSpaceRajan]

Thanks! I guess the only test failing is having more than 5 commits in the merge. I think you should be able to squash and merge and then it should not be a problem, right?

[IvyZX]

Thanks! I guess the only test failing is having more than 5 commits in the merge. I think you should be able to squash and merge and then it should not be a problem, right?

Sorry - we cannot squash it during merge because merging will be done by an automated system. Can you please squash them? Then I can proceed with merging.

[RaghuSpaceRajan]

Sorry, I tried to rebase it locally and squash the last 8 commits (as it shows my branch is 8 commits ahead on the GitHub UI), but locally it seems to show 26 commits since my first commit to the branch. I am pretty confused now. Would it be possible to squash and merge this one manually? Otherwise I'll create a new pull request with all changes made in a single commit manually. What do you think?