cuquantum.densitymat.jax.ElementaryOperator

class cuquantum.densitymat.jax.ElementaryOperator(data: jax.Array, callback: cuquantum.bindings.cudensitymat.WrappedTensorCallback | None = None, grad_callback: cuquantum.bindings.cudensitymat.WrappedTensorGradientCallback | None = None, offsets: Tuple[int, ...] = ())[source]

PyTree class for cuDensityMat’s elementary operator.

Methods

__init__(data: jax.Array, callback: cuquantum.bindings.cudensitymat.WrappedTensorCallback | None = None, grad_callback: cuquantum.bindings.cudensitymat.WrappedTensorGradientCallback | None = None, offsets: Tuple[int, ...] = ()) → None[source]

Initialize an ElementaryOperator object.

Parameters

  • data – Data buffer of the elementary operator.

  • callback – Forward callback for the elementary operator.

  • grad_callback – Gradient callback for the elementary operator.

  • offsets – Diagonal offsets of the elementary operator.