code: def init_step(self, x):
      rng = elegy.KeySeq(0)
      gx, g_params, g_states = self.generator.init(rng=rng)(x)
      dx, d_params, d_states = self.discriminator.init(rng=rng)(gx)
  
      g_optimizer_states = self.g_optimizer.init(g_params)
      d_optimizer_states = self.d_optimizer.init(d_params)
  
      return elegy.States(
          g_states=g_states,
          d_states=d_states,
          g_params=g_params,
          d_params=d_params,
          g_opt_states=g_optimizer_states,
          d_opt_states=d_optimizer_states,
          rng=rng,
          step=0,
      )
  
file_path: elegy/examples/need-fixing/WGAN-GP/model.py
imports: import jax
  import jax.numpy
  import optax
  import elegy
label_desc: Use image node_id 2 for calling the WGAN_GP obj's underlying member method code with example usage: obj.init_step(x) and returns: elegy