code: def get_batch_fn(ds: tp.Any) -> tp.Callable:
      """Returns either the batch_fn of the argument if it has one, otherwise `default_batch_fn`
      to allow arrays or datasets that don't inherit from elegy.data.Dataset
      """
      return getattr(ds, "batch_fn", default_batch_fn)
  
imports: import multiprocessing.pool
  import os
  import typing
  import jax
  import jax.numpy
  import numpy
  from .data_adapter import DataAdapter
  from .utils import is_none_or_empty
label_desc: Use image node_id 2 for calling a global function with example usage: get_batch_fn(ds) and returns: getattr