code: def init_step(self: M, key: jnp.ndarray, inputs: jnp.ndarray) -> M:
      features_in = np.prod(inputs.shape[1:])
  
      self.w = jax.random.uniform(
          key,
          shape=[
              features_in,
              self.features_out,
          ],
      )
      self.b = jnp.zeros([self.features_out])
  
      assert self.optimizer is not None
      self.optimizer = self.optimizer.init(self)
  
      return self
  
file_path: elegy/examples/jax/linear_classifier_train_step.py
imports: import os
  import typing
  from datetime import datetime
  import jax
  import jax.numpy
  import numpy
  import optax
  import typer
  from datasets.load import load_dataset
  import elegy
label_desc: Use image node_id 2 for calling the Model obj's underlying member method code with example usage: obj.init_step(key, inputs) and returns: self