Using pmap with a sparse matrix #13930
-
Hey, I'm trying to use a batched COO sparse matrix with
This fails with the following error message:
Instead, If I run the following:
I get this other error message:
Am I doing anything wrong? Does I'm using the following jax and jaxlib versions
|
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 2 replies
-
Thanks for the report – this looks like a bug. I'm not aware of anyone attempting to use |
Beta Was this translation helpful? Give feedback.
Thanks for the report – this looks like a bug. I'm not aware of anyone attempting to use
jax.experimental.sparse
withpmap
, so I suspect this is simply not possible in the current version of JAX, but I'll take a look to see if I can fix the issue.