neuralqx.vqs package

This module implements the variational quantum state backbone of neuraLQX.

class VariationalState(hilbert)

Bases: ABC

Abstract class for variational states representing either pure states or mixed quantum states.

A variational state is a quantum state depending on a set of parameters, and that supports operations such as computing quantum expectation values and their gradients.

A Variational stat can be serialized using flax’s msgpack machinery. See their docs.

expect(O)

Estimates the quantum expectation value for a given operator \(O\) or generic observable. In the case of a pure state \(\psi\) and an operator, this is \(\langle O\rangle= \langle \Psi|O|\Psi\rangle/\langle\Psi|\Psi\rangle\) otherwise for a mixed state \(\rho\), this is \(\langle O\rangle= \textrm{Tr}[\rho \hat{O}]/\textrm{Tr}[\rho]\).

Parameters:

O (AbstractOperator) – the operator or observable for which to compute the expectation value.

Return type:

Stats

Returns:

An estimation of the quantum expectation value \(\langle O\rangle\).

expect_and_grad(O, *, mutable=None, **kwargs)

Estimates the quantum expectation value and its gradient for a given operator \(O\).

Parameters:

  • O (AbstractOperator) – The operator \(O\) for which expectation value and gradient are computed.

  • mutable (Union[bool, str, Collection[str], DenyList, None]) – Can be bool, str, or list. Specifies which collections in the model_state should be treated as mutable: bool: all/no collections are mutable. str: The name of a single mutable collection. list: A list of names of mutable collections. This is used to mutate the state of the model while you train it (for example to implement BatchNorm. Consult Flax’s Module.apply documentation for a more in-depth explanation).

  • use_covariance – whether to use the covariance formula, usually reserved for hermitian operators, \(\textrm{Cov}[\partial\log\psi, O_{\textrm{loc}}\rangle]\)

Return type:

tuple[Stats, Any]

Returns:

An estimate of the quantum expectation value <O>. An estimate of the gradient of the quantum expectation value <O>.

expect_and_forces(O, *, mutable=None)

Estimates the quantum expectation value and the corresponding force vector for a given operator O.

The force vector \(F_j\) is defined as the covariance of log-derivative of the trial wave function and the local estimators of the operator. For complex holomorphic states, this is equivalent to the expectation gradient \(\frac{\partial\langle O\rangle}{\partial(\theta_j)^\star} = F_j\). For real-parameter states, the gradient is given by \(\frac{\partial\partial_j\langle O\rangle}{\partial\partial_j\theta_j} = 2 \textrm{Re}[F_j]\).

Parameters:

  • O (AbstractOperator) – The operator O for which expectation value and force are computed.

  • mutable (Union[bool, str, Collection[str], DenyList, None]) –

    Can be bool, str, or list. Specifies which collections in the model_state should be treated as mutable: bool: all/no collections are mutable. str: The name of a single mutable collection. list: A list of names of mutable collections. This is used to mutate the state of the model while you train it (for example to implement BatchNorm. Consult Flax’s Module.apply documentation for a more in-depth explanation).

Return type:

tuple[Stats, Any]

Returns:

An estimate of the quantum expectation value <O>. An estimate of the force vector \(F_j = \textrm{Cov}[\partial_j\log\psi, O_{\textrm{loc}}]\).

class VariationalMixedState(hilbert, *args, **kwargs)

Bases: VariationalState

expect(vstate, operator)

Computes the expectation value of the given operator over the variational state.

Additional Information:

To implement vstate.expect for a custom operator, implement the multiple-dispatch (plum-dispatc) based method according

Parameters:
Returns:

The expectation value wrapped in a Stats object.

expect(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator, chunk_size: NoneType) → netket.stats.mc_stats.Stats

# Standard implementation of expect for an MCState (pure) and a generic operator # The dispatch rule is not strictly needed, as everything currently implemented # in NetKet only defines a custom get_local_kernel_arguments and get_local_kernel # but if somebody wants to override behaviour for an existing operator or define # a completely arbitrary novel type of operator, this makes it much easier.

expect(vstate: netket.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_observable.AbstractObservable)
expect(vstate: netket.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_observable.AbstractObservable, chunk_size: int | tuple)
expect(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator, chunk_size: int) → netket.stats.mc_stats.Stats
expect(vstate: netket.vqs.full_summ.state.FullSumState, Ô: netket.operator._discrete_operator.DiscreteOperator) → netket.stats.mc_stats.Stats
expect(vstate: netket.vqs.mc.mc_state.state.MCState, op: netket.experimental.observable.renyi2.S2_operator.Renyi2EntanglementEntropy, chunk_size: int | None)
expect(vstate: netket.vqs.full_summ.state.FullSumState, op: netket.experimental.observable.renyi2.S2_operator.Renyi2EntanglementEntropy)
expect(vstate: netket.vqs.mc.mc_state.state.MCState, variance_operator: netket.experimental.observable.variance.variance_operator.VarianceObservable, chunk_size: NoneType)
expect(vstate: netket.vqs.full_summ.state.FullSumState, variance_operator: netket.experimental.observable.variance.variance_operator.VarianceObservable)
expect(vstate: netket.vqs.mc.mc_state.state.MCState, op: netket.experimental.observable.infidelity.infidelity_operator.InfidelityOperator, chunk_size: NoneType)
expect(vstate: netket.vqs.full_summ.state.FullSumState, op: netket.experimental.observable.infidelity.infidelity_operator.InfidelityOperator)
expect(vstate: netket._src.vqs.fermion_mf.state.DeterminantVariationalState, operator: netket.operator._fermion2nd.jax.FermionOperator2ndJax)
expect(vstate: netket._src.vqs.fermion_mf.state.DeterminantVariationalState, operator: netket._src.operator.particle_number_conserving_fermionic.operators.ParticleNumberConservingFermioperator2nd | netket._src.operator.particle_number_conserving_fermionic.operators.ParticleNumberAndSpinConservingFermioperator2nd)
expect(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, chunk_size: NoneType) → netket.stats.mc_stats.Stats
expect(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: Sequence[netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable], chunk_size: NoneType) → netket.stats.mc_stats.Stats

Extends the functionality of NetKet’s expect() to handle a sequence of operators [Ô_1, …, Ô_N]. Essentially, this computes <Ô_1> + … + <Ô_N> instead of <Ô_1 + … + Ô_N>.

expect(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_observable.AbstractObservable)
expect(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, chunk_size: int | tuple)
expect(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: Sequence[netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable], chunk_size: int | tuple)
expect(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, chunk_size: int) → netket.stats.mc_stats.Stats
expect(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: Sequence[netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable], chunk_size: int) → netket.stats.mc_stats.Stats

Extends the functionality of NetKet’s expect() to handle a sequence of operators [Ô_1, …, Ô_N]. Essentially, this computes <Ô_1> + … + <Ô_N> instead of <Ô_1 + … + Ô_N>.

expect_and_grad(vstate, operator, *args, mutable, **kwargs)

Estimates the quantum expectation value and its gradient for a given operator O.

See VariationalState.expect_and_grad docstring for more information.

Additional Information:

To implement vstate.expect for a custom operator, implement the multiple-dispatch (plum-dispatc) based method according to the signature below.

expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator, chunk_size: int | None, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared, chunk_size: int | None, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_observable.AbstractObservable, **kwargs)
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_observable.AbstractObservable, chunk_size: int | tuple, *args, **kwargs)
expect_and_grad(vstate: netket.vqs.full_summ.state.FullSumState, Ô: netket.operator._discrete_operator.DiscreteOperator, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None, **kwargs) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState, variance_operator: netket.experimental.observable.variance.variance_operator.VarianceObservable, chunk_size: NoneType, *, mutable)
expect_and_grad(vstate: netket.vqs.full_summ.state.FullSumState, variance_operator: netket.experimental.observable.variance.variance_operator.VarianceObservable, *, mutable)
expect_and_grad(vstate: netket.vqs.full_summ.state.FullSumState, op: netket.experimental.observable.infidelity.infidelity_operator.InfidelityOperator, *, mutable)
expect_and_grad(vstate: netket._src.vqs.fermion_mf.state.DeterminantVariationalState, operator: netket.operator._fermion2nd.jax.FermionOperator2ndJax, *args, use_covariance: bool | None = None, mutable, **kwargs)
expect_and_grad(vstate: netket._src.vqs.fermion_mf.state.DeterminantVariationalState, operator: netket._src.operator.particle_number_conserving_fermionic.operators.ParticleNumberConservingFermioperator2nd | netket._src.operator.particle_number_conserving_fermionic.operators.ParticleNumberAndSpinConservingFermioperator2nd, *args, use_covariance: bool | None = None, mutable, **kwargs)
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.experimental.observable.variance.variance_operator.VarianceObservable, chunk_size: NoneType, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance=None)

Variance cost specific path for the gradient

expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: Sequence[netket.operator._abstract_operator.AbstractOperator], chunk_size: int | None, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None) → tuple[netket.stats.mc_stats.Stats, Any]

Extended to handle multiple operators. If use_covariance=True, we use the new expect_and_forces() that can handle a list.

expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator, chunk_size: int | None, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared, chunk_size: int | None, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: Sequence[netket.operator._lazy.Squared | neuralqx.operators._lazy.PenaltyCost | neuralqx.operators.types.computational_operator.base.ComputationalOperator], chunk_size: int | None, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None) → tuple[netket.stats.mc_stats.Stats, Any]

Extends the squared-op code to handle either a single Squared[…] operator, or a list of Squared[…] operators. If it’s a list, we do a multi-operator approach.

It can also accept PenaltyCost operators which are not Squared.

expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators._lazy.InverseExpectationCost, chunk_size: int | None, *args, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False, use_covariance: bool | None = None) → tuple[netket.stats.mc_stats.Stats, Any]
IEC-specific, computes gradient via forces on E only:

d/dθ <IEC> = f’(<E>) * d/dθ <E>

while keeping the forward stats of IEC via the (already implemented) local estimator

Note: the assumption here is that the wrapped operator E is Hermitian. This will produce incorrect results if it is not, as we use the covariance method to compute the gradient

expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, **kwargs)
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: Sequence[netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable], **kwargs)
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, chunk_size: int | tuple, *args, **kwargs)
expect_and_grad(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: Sequence[netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable], chunk_size: int | tuple, *args, **kwargs)
expect_and_forces(vstate, operator, *args, mutable, **kwargs)

Estimates the quantum expectation value and corresponding force vector for a given operator O.

See VariationalState.expect_and_forces docstring for more information.

Additional Information:

To implement vstate.expect for a custom operator, implement the multiple-dispatch (plum-dispatc) based method according to the signature below.

expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator, chunk_size: NoneType, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_observable.AbstractObservable, **kwargs)
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_observable.AbstractObservable, chunk_size: Any, *args, **kwargs)
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator, chunk_size: int, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_forces(vstate: netket.vqs.full_summ.state.FullSumState, Ô: netket.operator._discrete_operator.DiscreteOperator, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_forces(vstate: netket._src.vqs.fermion_mf.state.DeterminantVariationalState, operator: netket.operator._fermion2nd.jax.FermionOperator2ndJax, *args, mutable, **kwargs)
expect_and_forces(vstate: netket._src.vqs.fermion_mf.state.DeterminantVariationalState, operator: netket._src.operator.particle_number_conserving_fermionic.operators.ParticleNumberConservingFermioperator2nd | netket._src.operator.particle_number_conserving_fermionic.operators.ParticleNumberAndSpinConservingFermioperator2nd, *args, mutable, **kwargs)
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: Sequence[netket.operator._abstract_operator.AbstractOperator], chunk_size: NoneType, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False) → tuple[netket.stats.mc_stats.Stats, Any]

Extends the NetKet logic of expect_and_forces() to apply to a sequence of operators.

expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator, chunk_size: NoneType, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, **kwargs)
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, chunk_size: Any, *args, **kwargs)
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, operator: Sequence[netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable], chunk_size: Any, *args, **kwargs)
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable, chunk_size: int, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False) → tuple[netket.stats.mc_stats.Stats, Any]
expect_and_forces(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: Sequence[netket.operator._abstract_operator.AbstractOperator | netket.operator._abstract_observable.AbstractObservable], chunk_size: int, *, mutable: bool | str | Collection[str] | flax.core.scope.DenyList = False) → tuple[netket.stats.mc_stats.Stats, Any]

Extends the NetKet logic of expect_and_forces() to apply to a sequence of operators.

class MCState(sampler, model=None, *, n_samples=None, n_samples_per_rank=None, n_discard_per_chain=None, chunk_size=None, variables=None, init_fun=None, apply_fun=None, seed=None, sampler_seed=None, mutable=False, training_kwargs={}, is_group_averaged=False)

Bases: VariationalState

Variational State for a Variational Neural Quantum State.

The state is sampled according to the provided sampler.

expect(O)

Estimates the quantum expectation value for a given operator \(O\) or generic observable. In the case of a pure state \(\psi\) and an operator, this is \(\langle O\rangle= \langle \Psi|O|\Psi\rangle/\langle\Psi|\Psi\rangle\) otherwise for a mixed state \(\rho\), this is \(\langle O\rangle= \textrm{Tr}[\rho \hat{O}]/\textrm{Tr}[\rho]\).

Parameters:

O (Union[AbstractOperator, Sequence[AbstractOperator]]) – the operator or observable for which to compute the expectation value.

Return type:

Stats

Returns:

An estimation of the quantum expectation value \(\langle O\rangle\).

expect_and_grad(O, *, mutable=None, **kwargs)

Estimates the quantum expectation value and its gradient for a given operator \(O\).

Parameters:

  • O (Union[AbstractOperator, Sequence[AbstractOperator]]) – The operator \(O\) for which expectation value and gradient are computed.

  • mutable (Union[bool, str, Collection[str], DenyList, None]) –

    Can be bool, str, or list. Specifies which collections in the model_state should be treated as mutable: bool: all/no collections are mutable. str: The name of a single mutable collection. list: A list of names of mutable collections. This is used to mutate the state of the model while you train it (for example to implement BatchNorm. Consult Flax’s Module.apply documentation for a more in-depth explanation).

  • use_covariance – whether to use the covariance formula, usually reserved for hermitian operators, \(\textrm{Cov}[\partial\log\psi, O_{\textrm{loc}}\rangle]\)

Return type:

tuple[Stats, Any]

Returns:

An estimate of the quantum expectation value <O>. An estimate of the gradient of the quantum expectation value <O>.

expect_and_forces(O, *, mutable=None)

Estimates the quantum expectation value and the corresponding force vector for a given operator O.

The force vector \(F_j\) is defined as the covariance of log-derivative of the trial wave function and the local estimators of the operator. For complex holomorphic states, this is equivalent to the expectation gradient \(\frac{\partial\langle O\rangle}{\partial(\theta_j)^\star} = F_j\). For real-parameter states, the gradient is given by \(\frac{\partial\partial_j\langle O\rangle}{\partial\partial_j\theta_j} = 2 \textrm{Re}[F_j]\).

Parameters:

  • O (Union[AbstractOperator, Sequence[AbstractOperator]]) – The operator O for which expectation value and force are computed.

  • mutable (Union[bool, str, Collection[str], DenyList, None]) –

    Can be bool, str, or list. Specifies which collections in the model_state should be treated as mutable: bool: all/no collections are mutable. str: The name of a single mutable collection. list: A list of names of mutable collections. This is used to mutate the state of the model while you train it (for example to implement BatchNorm. Consult Flax’s Module.apply documentation for a more in-depth explanation).

Return type:

tuple[Stats, Any]

Returns:

An estimate of the quantum expectation value <O>. An estimate of the force vector \(F_j = \textrm{Cov}[\partial_j\log\psi, O_{\textrm{loc}}]\).

get_local_kernel_arguments(vstate, Ô)

Returns the samples of vstate used to compute the expectation value of the operator O, and the connected elements and matrix elements.

Parameters:

  • vstate (Any) – the variational state

  • Ô – the operator

Returns:

A Tuple with 2 elements (sigma, args), where the first elements should be the samples over which the classical expectation value should be computed, while the latter is anything that can be fed as input to the local_kernel.

get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator.DiscreteOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._continuous_operator.ContinuousOperator)
get_local_kernel_arguments(vs: netket.vqs.mc.mc_state.state.MCState, O: netket.operator._sum.operator.SumGenericOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._discrete_operator.DiscreteOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._abstract_super_operator.AbstractSuperOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._lazy.Squared[netket.operator._abstract_super_operator.AbstractSuperOperator])
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators._lazy.PenaltyCost)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator.DiscreteOperator | neuralqx.operators.types._discrete_operator.DiscreteOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._continuous_operator.ContinuousOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.experimental.observable.variance.variance_operator.VarianceObservable)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators.types.computational_operator.base.ComputationalOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators.types.computational_operator.jax.ComputationalJaxOperator)
get_local_kernel_arguments(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators._lazy.InverseExpectationCost)

Build the args for the InverseExpectationCost with batch-consistent (σ, inner_args)

We will:

  • collapse σ like _expect() does

  • build inner_args from that collapsed σ

  • compute <V> from the same pair

  • pack dynamic scalars (fprime, g)

  • return the σ_collapsed so later kernels see the same layout and avoid recompilation

get_local_kernel(vstate, Ô)

Returns the function computing the local estimator for the given variational state and operator.

Parameters:

  • vstate (Any) – the variational state

  • Ô – the operator

Returns:

A callable accepting the output of get_configs(vstate, O).

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator.DiscreteOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._continuous_operator.ContinuousOperator)
get_local_kernel(vs: netket.vqs.mc.mc_state.state.MCState, O: netket.operator._sum.operator.SumGenericOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared, chunk_size: int)

# Dispatches to select what expect-kernel to use

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator, chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator.DiscreteOperator, chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState, Ô: netket.operator._continuous_operator.ContinuousOperator, chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._abstract_super_operator.AbstractSuperOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._discrete_operator.DiscreteOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._lazy.Squared[netket.operator._abstract_super_operator.AbstractSuperOperator])
get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._abstract_super_operator.AbstractSuperOperator, chunk_size: int)

# Dispatches to select what expect-kernel to use

get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._lazy.Squared[netket.operator._abstract_super_operator.AbstractSuperOperator], chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._discrete_operator.DiscreteOperator, chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_mixed_state.state.MCMixedState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator, chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators._lazy.PenaltyCost)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator.DiscreteOperator | neuralqx.operators.types._discrete_operator.DiscreteOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._continuous_operator.ContinuousOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.experimental.observable.variance.variance_operator.VarianceObservable)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators.types.computational_operator.base.ComputationalOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators.types.computational_operator.jax.ComputationalJaxOperator)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators._lazy.InverseExpectationCost)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator.DiscreteOperator | neuralqx.operators.types._discrete_operator.DiscreteOperator, chunk_size: int)

# standard numba operators

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._discrete_operator_jax.DiscreteJaxOperator, chunk_size: int)

# standard JAX operators

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._continuous_operator.ContinuousOperator, chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.operator._lazy.Squared, chunk_size: int)

# Squared operators

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators.types.computational_operator.base.ComputationalOperator, chunk_size: int)

# standard computational operators

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators.types.computational_operator.jax.ComputationalJaxOperator, chunk_size: int)

# standard computational JAX operators

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators._lazy.PenaltyCost, chunk_size: int)

# standard PenaltyCost operators

get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: neuralqx.operators._lazy.InverseExpectationCost, chunk_size: int)
get_local_kernel(vstate: netket.vqs.mc.mc_state.state.MCState | neuralqx.vqs.mc.mc_state.state.MCState, Ô: netket.experimental.observable.variance.variance_operator.VarianceObservable, chunk_size: int)

# VarianceObservable not supported for now

Subpackages

Submodules