code: def _partial_flatten(x):
      """Flatten all but the first dimension of an ndarray."""
      return np.reshape(x, (x.shape[0], -1))
  
file_path: jax/examples/datasets.py
imports: import array
  import gzip
  import os
  from os import path
  import struct
  import urllib.request
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: _partial_flatten(x) and returns: np