How to use init_fun and apply_fun in MCState

[whatisspin]

Hello,

I am trying to add a dropout layer, for which I need to pass the rng to the apply function of the model. But I cannot figure out how to pass the init_fun and apply_fun to the MCState. For example (even without the dropout):

import netket as nk
import jax
import jax.nn as jnn
import flax.linen as nn
import jax.numpy as jnp
import numpy as np
import optax


L=10

hi= nk.hilbert.Spin(s=0.5,N=L)
g= nk.graph.Chain(length=10)
ham= nk.operator.Heisenberg(hilbert=hi, graph=g)
sa= nk.sampler.MetropolisExchange(hilbert=hi,graph=g)

class FFNN(nn.Module):
    alpha: int
    
    @nn.compact
    def __call__(self,x):
        #some_param= self.param('some_param',nn.initializers.zeros_init(),(1,))
        #dropout_rng= self.make_rng('dropout')
        
        x= nn.Dense(features= self.alpha*L, param_dtype=np.complex128, kernel_init= jnn.initializers.normal(0.01), bias_init=jnn.initializers.normal(stddev=0.01), name="layer1")(x)
        x= nk.nn.reim_relu(x)       
        
        #x= nn.Dropout(0.5, deterministic= False)(x)
        
        x= jnp.sum(x, axis=(-1))
        return x
model= FFNN(alpha=1)

key1, key2= jax.random.split(jax.random.PRNGKey(0),2)
x= hi.random_state(key=key1,size=2*11)
params= model.init(rngs=key2,x=x)
y= model.apply(variables=params,x=x)

#vs= nk.vqs.MCState(sa,model=model,n_samples=2**11)
vs= nk.vqs.MCState(sampler=sa,init_fun=model.init(rngs=key2,x=x),apply_fun=model.apply(variables=params,x=x),n_samples=2**11)

sr= nk.optimizer.SR(diag_shift=0.1)
gs= nk.driver.VMC(ham,optimizer=optax.sgd(0.05),variational_state=vs,preconditioner=sr)

gs.run(1000)

I have the following error:

Click me

> Traceback (most recent call last):
  File "temp_delete.py", line 44, in <module>
    vs= nk.vqs.MCState(sa,init_fun=model.init(rngs=key2,x=x),apply_fun=model.apply(variables=params,x=x),n_samples=2**11)
  File ".../lib/python3.8/site-packages/netket/vqs/mc/mc_state/state.py", line 240, in __init__
    self.init(seed, dtype=sampler.dtype)
  File ".../lib/python3.8/site-packages/netket/vqs/mc/mc_state/state.py", line 275, in init
    variables = jit_evaluate(self._init_fun, {"params": key}, dummy_input)
ValueError: Non-hashable static arguments are not supported. An error occurred during a call to 'jit_evaluate' while trying to hash an object of type <class 'flax.core.frozen_dict.FrozenDict'>, FrozenDict({
    params: {
        layer1: {
            kernel: Array([[-0.01087773-1.20742927e-02j,  0.01457856+4.03403656e-05j,
                    -0.0011959 +6.51546689e-03j, -0.00431869+7.91588313e-03j,
                    -0.01119718+4.63213851e-03j, -0.00697697+8.61602995e-03j,
                    -0.00492156-5.59540416e-03j, -0.00312104-2.63089893e-03j,
                    -0.00126524-2.85040372e-03j,  0.0013996 -2.55798337e-03j],
                   [-0.0038206 +3.42691290e-03j, -0.0071288 -6.04020543e-04j,
                     0.0042634 -4.88060986e-03j,  0.00615774-1.18782339e-02j,
                     0.00085392-1.19015412e-02j,  0.00876033+6.16646076e-03j,
                     0.00617935-1.20160263e-03j, -0.00962342+5.27570617e-03j,
                    -0.00793193+9.30064053e-03j, -0.00189154-5.69452869e-03j],
                   [ 0.00736881-5.92196217e-04j,  0.00517723-9.09916638e-03j,
                     0.00354315+2.98299989e-03j,  0.00431793+2.66163265e-03j,
                     0.0063637 +1.78909313e-03j, -0.00429622-1.40461298e-03j,
                     0.0011257 +1.31498433e-02j,  0.00246002+5.50134297e-03j,
                     0.00551634+5.07662086e-03j,  0.00737985+7.19355200e-03j],
                   [-0.0055389 +6.41792032e-04j,  0.0034731 +3.05790167e-03j,
                    -0.00610383+1.56022457e-02j, -0.01035439-1.66959746e-02j,
                     0.00193768+2.03921050e-03j,  0.00257564-2.95988938e-03j,
                     0.00263892-6.12085503e-03j, -0.0050302 -1.05616391e-02j,
                    -0.0085371 +2.94148657e-03j, -0.00044882+2.99630965e-03j],
                   [-0.00283266-1.09127341e-02j,  0.00320202+3.98983813e-03j,
                     0.00411279-9.88615510e-04j, -0.00215587-2.90032772e-03j,
                    -0.00496469+1.24980484e-02j,  0.00128656+5.48089969e-03j,
                    -0.0060321 +6.34088432e-03j, -0.00596992-1.55176917e-04j,
                     0.00871737+7.96754836e-03j,  0.00491199+2.57438968e-04j],
                   [-0.0009041 -1.05545951e-02j, -0.00248048-3.87302547e-03j,
                    -0.00331804-1.09678106e-03j,  0.01169915-1.51457821e-03j,
                     0.00071176-1.42297009e-02j, -0.00599372-3.79475626e-03j,
                    -0.01129931-3.10419746e-03j,  0.00117312+1.86229295e-03j,
                     0.00211868+8.90366332e-03j, -0.00188821-2.91471704e-03j],
                   [-0.01622197-1.08967473e-02j,  0.00266872+3.00000357e-03j,
                     0.01320552+2.62331463e-03j,  0.00542052-1.01636301e-02j,
                     0.01542558+6.01395042e-03j,  0.00122886-6.43474437e-03j,
                     0.00146662+5.00728180e-03j,  0.00512276+1.58961343e-04j,
                     0.01555954-8.55212073e-04j, -0.00334311+1.67141399e-03j],
                   [ 0.00754611-7.46525690e-03j, -0.00090699-2.22644000e-03j,
                     0.00061078-1.22834251e-02j, -0.00669451-4.56227107e-03j,
                     0.01211781-7.46405307e-03j,  0.00446219+9.33810820e-03j,
                     0.0063756 -1.13825486e-02j, -0.00026557+6.09254106e-03j,
                    -0.00842443+3.55875895e-03j, -0.00631857-5.68531050e-03j],
                   [-0.01576199+6.12012151e-03j,  0.01748999-8.00780006e-03j,
                    -0.00237387+2.71730781e-03j,  0.00165019+9.60639262e-04j,
                    -0.00903517+2.09651703e-03j, -0.00288515+4.75256797e-04j,
                    -0.00428972-1.69981993e-04j, -0.003133  -7.94254698e-04j,
                     0.00042122-1.47978491e-03j, -0.00049524+4.94540996e-03j],
                   [ 0.00595958-3.65619226e-03j, -0.00658993+9.05791103e-03j,
                    -0.01348982-3.34409723e-03j,  0.00163019+3.00761646e-03j,
                    -0.01188033+1.22092102e-02j, -0.00853378+2.02675056e-03j,
                    -0.00143499+8.68940789e-03j,  0.01017866+5.97673398e-03j,
                    -0.00211163-6.24396033e-03j,  0.00413886+8.31357373e-03j]],      dtype=complex128),
            bias: Array([-0.00226185+0.00224257j, -0.00187046-0.00376108j,
                    0.00958692+0.00215852j,  0.00172275-0.00031328j,
                   -0.0032391 -0.0076107j ,  0.00225818-0.00716646j,
                    0.00029576-0.00637876j, -0.01035107+0.00064331j,
                   -0.00013353-0.01016284j,  0.00498229+0.00611868j],      dtype=complex128),
        },
    },
}). The error was:
TypeError: unhashable type: 'ArrayImpl'

While I do get the output for y=model.apply(variables=params,x=x), what would be the correct way to pass these functions to MCState?

Thank you.

[PhilipVinc]

So a few things:

• init_fun and apply_fun must be functions. You are passing the result of calling this function which is a pytree. You should either pass init_fun=model.init and apply_fun=model.apply or a partial covering those.

• The changes to make netket work with dropout layers (essentially, passing those extra fields) are not in netket, but I had designed it in such a way to allow to insert them at some point. So, netket will not easily work with dropout right now but I can try to implement it in the next few days.

For that to happen, however, you should provide me some code that works to initialise the network with the dropout layer and how to evaluate it during/not during training (With what kwargs to model.apply)

[whatisspin]

Thank you for the reply.

So maybe a code like this?

import netket as nk
import jax
import jax.nn as jnn
import flax.linen as nn
import jax.numpy as jnp
import numpy as np

L=10
hi= nk.hilbert.Spin(s=0.5,N=L)

class FFNN(nn.Module):
    
    alpha: int
    
    @nn.compact
    def __call__(self,x,training:bool):
        some_param= self.param('some_param',nn.initializers.zeros_init(),(1,))
        dropout_rng= self.make_rng('dropout')
        
        x= nn.Dense(features= self.alpha*L, param_dtype=np.complex128, kernel_init= jnn.initializers.normal(0.01), bias_init=jnn.initializers.normal(stddev=0.01), name="layer1")(x)
        x= nk.nn.reim_relu(x)       
        
        x= nn.Dropout(0.5, deterministic= not training)(x)

        x= nn.Dense(features= self.alpha*L, param_dtype=np.complex128, kernel_init= jnn.initializers.normal(0.01), bias_init=jnn.initializers.normal(stddev=0.01), name="layer2")(x)
        x= nk.nn.reim_relu(x)
        
        x= jnp.sum(x, axis=(-1))
        return x
model= FFNN(alpha=1)

x= hi.random_state(key=jax.random.PRNGKey(1), size=2**11)
params= model.init(rngs={'params':jax.random.PRNGKey(0), 'dropout':jax.random.PRNGKey(2)}, x=x, training=False)

y= model.apply(rngs={'dropout': jax.random.PRNGKey(5)},variables=params, x=x, training=True)