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 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

Contents

Using PennyLane

Release news

Development

API

Internals