code: def train_validation_split(
      arrays, validation_split, shuffle=True
  ) -> tp.Tuple:
      """
      Split arrays into random train and validation subsets.
      Arguments:
          arrays: Arrays to split. Allowed inputs are arbitrarily nested structures
              of Jax and NumPy arrays.
          validation_split: Float between 0 and 1. The proportion of the dataset to
              include in the validation split. The rest of the dataset will be included
              in the training split.
          shuffle: Bool. Whether to shuffle the data before performing a split. If
              `False`, the last `validation_split` fraction of that training data will
              become the validation split.
      Returns:
          `(train_arrays, validation_arrays)`
      """
  
      def _can_split(t):
          supported_types = (jnp.ndarray, np.ndarray)
          # if pd:
          #     supported_types = (jnp.ndarray, np.ndarray, pd.Series, pd.DataFrame)
          return isinstance(t, supported_types) or t is None
  
      flat_arrays = flatten(arrays)
      # flat_arrays = arrays
      if not all(_can_split(t) for t in arrays):
          raise ValueError(
              "`validation_split` is only supported for Tensors or NumPy "
              "arrays, found: {}".format(arrays)
          )
  
      if all(t is None for t in flat_arrays):
          return arrays, arrays
  
      first_non_none = None
      for t in flat_arrays:
          if t is not None:
              first_non_none = t
              break
  
      # Assumes all arrays have the same batch shape or are `None`.
      batch_dim = int(first_non_none.shape[0])
      indices = np.arange(batch_dim)
      if shuffle:
          indices = np.random.shuffle(indices)
      split_at = int(math.floor(batch_dim * (1.0 - validation_split)))
      train_indices = indices[:split_at]
      val_indices = indices[split_at:]
  
      def _split(t, indices):
          if t is None:
              return t
          return t[indices]
  
      train_arrays = map_structure(
          functools.partial(_split, indices=train_indices), arrays
      )
      val_arrays = map_structure(
          functools.partial(_split, indices=val_indices), arrays
      )
      return train_arrays, val_arrays
  
imports: import collections
  import functools
  import math
  import typing
  import jax
  import jax.numpy
  import numpy
  import six
  import typing_extensions
  from elegy import types, utils
label_desc: Use image node_id 10 for calling a global function with example usage: train_validation_split(arrays, validation_split, shuffle) and returns: train_arrays,  val_arrays, unknown, arrays,  arrays, t, t