code: def get_train_features():
      """
      Create features for the entire datasets.
      """
      global xs_train, ys_train
      print(x_train.shape)
      original_train_ = np.array(original_train)
      encoded_train_ = np.array(encoded_train)
  
      print("Computing features")
      ys_train = features(encoded_train_, True)
  
      patch_size = 16
      ss = original_train_.shape[3] // patch_size
      # Okay so this is an ugly transpose block.
      # We are going from [outer_batch, batch_size, channels, width, height
      # to [outer_batch, batch_size, channels, width/patch_size, patch_size, height/patch_size, patch_size]
      # Then we reshape this and flatten so that we end up with
      # [other_batch, batch_size, width/patch_size, height_patch_size, patch_size**2*channels]
      # So that now we can run features on the last dimension
      original_train_ = (
          original_train_.reshape(
              (
                  original_train_.shape[0],
                  original_train_.shape[1],
                  original_train_.shape[2],
                  ss,
                  patch_size,
                  ss,
                  patch_size,
              )
          )
          .transpose((0, 1, 3, 5, 2, 4, 6))
          .reshape(
              (
                  original_train_.shape[0],
                  original_train_.shape[1],
                  ss**2,
                  patch_size**2,
              )
          )
      )
  
      xs_train = features(original_train_, False)
  
      print(xs_train.shape, ys_train.shape)
  
imports: import collections
  import hashlib
  import time
  import multiprocessing
  import torch
  import numpy
  import torch.nn
  import scipy.stats
  import matplotlib.pyplot
  from PIL import Image
  import jax
  import jax.numpy
  import objax
  import scipy.optimize
  import numpy
  import multiprocessing
label_desc: Use image node_id 5 for calling a global function with example usage: get_train_features() without return types