Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve error message when user mistakenly holds a jax Array in an nnx.Module #4480

Open
RaghuSpaceRajan opened this issue Jan 13, 2025 · 2 comments · May be fixed by #4492
Open

Improve error message when user mistakenly holds a jax Array in an nnx.Module #4480

RaghuSpaceRajan opened this issue Jan 13, 2025 · 2 comments · May be fixed by #4492

Comments

@RaghuSpaceRajan
Copy link

The line here raises, imo, a vague error message (at least I felt so as a beginner to Jax, Flax and nnx).

I was holding a jax Array (a random key) in an nnx.Module and that led to the error. tbh, I am not sure if this happens only in this case. Either way, I would be happy to help improve the error message if it helps.
@cgarciae
Copy link
Collaborator

Hi @RaghuSpaceRajan, feel free to send a PR!

@RaghuSpaceRajan RaghuSpaceRajan linked a pull request Jan 20, 2025 that will close this issue
Open
@RaghuSpaceRajan
Copy link
Author

Hi @cgarciae , thanks. I made it more verbose. Does it look okay. It's in PR #4492.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
2 participants
@cgarciae @RaghuSpaceRajan