qml.capture.eval_jaxpr¶
- eval_jaxpr(jaxpr, consts, *args)[source]¶
A version of
jax.core.eval_jaxpr
that can handle creating arrays with dynamic shapes.- Parameters
jaxpr (jax.core.Jaxpr) – a jaxpr
consts (list[TensorLike]) – the constants for the jaxpr
*args (TensorLike) – the arguments for the jaxpr
- Returns
list[TensorLike]
This function only differs from
jax.core.eval_jaxpr
in that it can handle the creation of dynamically shaped arrays viaiota
andbroadcast_in_dim
. >>> import jax >>> jax.config.update(“jax_dynamic_shapes”, True) >>> def f(i): … return jax.numpy.arange(i) >>> jaxpr = jax.make_jaxpr(f)(3) >>> qml.capture.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 2) [Array([0, 1], dtype=int32)] >>>> jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 2) XlaRuntimeError: error: ‘mhlo.dynamic_iota’ op can’t be translated to XLA HLO