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]

Suggested change

	f'{type(value)}. Leaf values of this type are not supported for nnx.Modules.'
	f'{type(value)}. Leaf values of this type are not supported for nnx.Modules or nnx.Object in general.'

[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?