Closed pointeee closed 5 months ago
Hi Chenze, could you try ptcl_spacing = 1.0
instead and see if the error persists?
Yes, with ptcl_spacing = 1.0
it reports the same error ValueError: invalid literal for int() with base 10: 'int16[8,3]'
I don't know why that happens and the error message is vague.
For the current version, you don't need to jit because it's already done for the most part. But this may change in the future.
Thanks for the reply. I believe it is related to how Particles.__post_init__
in particles.py deals with non-tracing arrays when getting jitted.
But since jit is optional here, I think we can close the issue.
I believe it is related to how
Particles.__post_init__
in particles.py deals with non-tracing arrays when getting jitted.
FYI Particles
is instantiated in lpt()
which is already jitted
Hi, I am writing a sampling code using pmwd. The objective function in my code (actually negative log probability) is frequently called, so I attempted to add
@jax.jit
to speed it up. However, I always got the an error likeValueError: invalid literal for int() with base 10: 'int16[8,3]'
when executing thenbody
function. Here is an illustrative example:And I got the following error for the script above when enabling
@jax.jit
:I am using jax-0.4.23 [cpu] & pmwd-0.1.dev122+g1e1c634.d20240303 (GitHub commit 1e1c634). Is this an expected behavior?
Mar 6th Update: I attempt to run the script above in an environment with jax-gpu installation, and the error persists. I think it is not a CPU-version-only problem.