-
|
Really simple question, but I cannot find any answer to this question from the Jax documentation. If I decorate a function with 'shard_map', will it also be jitted, or do I need to decorate it with 'jit'? In which case, what is the most efficient order of composition order? Thanks! |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 1 reply
-
|
No, shard_map is eager by default. You will need to wrap a jit around it. |
Beta Was this translation helpful? Give feedback.
-
|
Yes that's what I thought. Thanks for the answer! |
Beta Was this translation helpful? Give feedback.
No, shard_map is eager by default. You will need to wrap a jit around it.
jit(shard_map)should be fine for your last question.