code: def init_model(layers=layers, rng=npr.RandomState(0)):
      model = []
      for (
          in_size,
          out_size,
      ) in zip(layers[:-1], layers[1:]):
          new_w = 0.1 * rng.randn(in_size, out_size)
          new_b = 0.1 * rng.randn(out_size)
          new_layer = (new_w, new_b)
          model.append(new_layer)
  
      return model
  
file_path: DALI/docs/examples/frameworks/jax/model.py
imports: from jax import jit, grad, lax, pmap
  from functools import partial
  import jax.numpy
  from jax.scipy.special import logsumexp
  import numpy.random
label_desc: Use image node_id 1 for calling a global function with example usage: init_model(layers, rng) and returns: model