How to add new variables to a Jaxpr? #28976
-
|
I would like to modify a Jaxpr in a way that I need to add additional variables to it. What would be the correct way to create new unique ids for these variables? I have seen some people determining the maximum id in the Jaxpr and then incrementing but that feels a bit hacky to me to be honest. |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 1 reply
-
|
YMMV because the Jaxpr data structure isn't really meant to be user-facing, but I don't believe you need to specify an id. Variables in a Jaxpr are new_var = Var("", ShapedArray((1, 5), jnp.float32))should be fine for creating a new variable! |
Beta Was this translation helpful? Give feedback.
-
|
Okay, thanks! I thought the ids were used for referencing internally somewhere. |
Beta Was this translation helpful? Give feedback.
YMMV because the Jaxpr data structure isn't really meant to be user-facing, but I don't believe you need to specify an id. Variables in a Jaxpr are
Varobjects and they only wrap aAbstractValue:should be fine for creating a new variable!