code: def test_state_returns_initial_value(self):
      with hk.variables():
          # NOTE: Initial value defined outside transform.
          v = tf.Variable(0, trainable=False)
  
      f = hk.transform_with_state(lambda: v.assign(1))
      params, state = f.init()
      initial_v = state[v.ref()]
      self.assertEqual(initial_v.numpy(), 0)
  
      y, state = f.apply(params, state)
      final_v = state[v.ref()]
      self.assertEqual(y.numpy(), 1)
      self.assertEqual(final_v.numpy(), 1)
  
imports: from absl.testing import parameterized
  import sonnet
  from sonnet.src import test_utils
  from sonnet.src.functional import haiku
  import tensorflow
  import tree
label_desc: Use image node_id 5 for calling the NetworkTest obj's underlying member method code with example usage: obj.test_state_returns_initial_value() without return types