code: def test_init(self):
      N = 0
  
      class Model(ModelCore):
          a: jnp.ndarray = tx.node()
  
          def __init__(self):
              super().__init__()
              self.a = jnp.array(-1, dtype=jnp.int32)
  
          def init_step(
              self,
              key: jnp.ndarray,
              inputs: tp.Any,
          ) -> "Model":
              nonlocal N
  
              N += 1
              self.a = jnp.array(0, dtype=jnp.int32)
              print("JITTING")
              return self
  
      model = Model()
      inputs = np.array(1.0)
  
      assert N == 0
      assert model.a == -1
  
      model.init_on_batch(inputs)
      assert N == 1
  
      # jits again because _initialized changed
      model.init_on_batch(inputs)
      assert N == 2
  
      # no jit change this time
      model.init_on_batch(inputs)
      assert N == 2
  
file_path: elegy/tests/model/model_core_test.py
imports: import typing
  import unittest
  import jax
  import jax.numpy
  import numpy
  import treex
  import elegy
  from elegy.model.model_core import ModelCore
label_desc: Use image node_id 1 for calling the ModelCoreTest obj's underlying member method code with example usage: obj.test_init() and returns: self