@partial function not working in custom operator definition

[singhavishek]

I am trying to understand custom operatir defining in Netket documentation given in this link
there I found @partial function in a cell. I am not able to get from which library this function is imported. I an getting error

`---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-10-4871a11b9c4d> in <module>()
     11   return eta, jnp.ones(1)
     12 
---> 13 @partial(jax.vmap, in_axes=(None, None, 0, 0, 0))
     14 def e_loc(logpsi, pars, sigma, eta, mels):
     15   return jnp.sum(mels * jnp.exp(logpsi(pars, eta) - logpsi(pars, eta)), axis=-1)

NameError: name 'partial' is not defined`

the cell is, 
`@jax.vmap
def get_conns_and_mels(sigma):
  # get number of spins
  N = sigma.shape[-1]
  # repeat eta N times
  eta = jnp.tile(sigma, (N,1))
  # diagonal indices
  ids = np.diag_indices(N)
  # flip those indices
  eta = eta.at[ids].set(-eta.at[ids].get())
  return eta, jnp.ones(1)

@partial(jax.vmap, in_axes=(None, None, 0,0,0))
def e_loc(logpsi, pars, sigma, eta, mels):
  return jnp.sum(mels * jnp.exp(logpsi(pars, eta) - logpsi(pars, eta)), axis=-1)`

[PhilipVinc]

from functools import partial

[PhilipVinc]

Yes but it is not as straightforward as for standard ground/excited states computations, and you must know a bit what you are doing in order to bolt together different parts.

You can either to it by computing the dynamics of the density matrix using the Lindblad and thermal dissipators (that's not ideal and will be computationally expensive, but will work) or you can perform imaginary time dynamics over a set of initial states.

[singhavishek]

what if, I define a custom operator for finite temperature expectation value and run it.

[PhilipVinc]

what is a 'finite temperature expectation' ? (I mean, there are many approaches that can work and you can implement them on top of NetKet. I am just unsure of what you are proposing)

[singhavishek]

may be here in chat I am not able to explain it properly. Can we discuss sometime over a video chat.
If possible.

[PhilipVinc]

If you want to chat you can join the slack where several of us hang out. The invite link is on a badge of https://github.com/netket/netket