code: def init_random_params(scale, layer_sizes, rng=npr.RandomState(0)):
      return [
          (scale * rng.randn(m, n), scale * rng.randn(n))
          for m, n, in zip(layer_sizes[:-1], layer_sizes[1:])
      ]
  
file_path: jax/examples/spmd_mnist_classifier_fromscratch.py
imports: from functools import partial
  import time
  import numpy
  import numpy.random
  import jax
  from jax import jit, grad, pmap
  from jax.scipy.special import logsumexp
  from jax.tree_util import tree_map
  from jax import lax
  import jax.numpy
  from examples import datasets
label_desc: Use image node_id 1 for calling a global function with example usage: init_random_params(scale, layer_sizes, rng) and returns: unknown