qml.capture.eval_jaxpr

eval_jaxpr(jaxpr, consts, *args)

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 via iota and broadcast_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

code/api/pennylane.capture.eval_jaxpr

 

Download Python script

 

Download Notebook

 

View on GitHub