def _length_of_batch(batch):
    """Get the number of samples in the mini-batch `batch`.

    `batch` can be:
    - a NumPy array, in which case `len(batch)` (size of first axis) will
      be returned
    - a tuple, in which case `_length_of_batch` will be invoked
      (recursively) on the first element

    As a consequence, mini-batches can be structured; lists and tuples can
    be nested arbitrarily deep.

    Parameters
    ----------
    batch: tuple or NumPy array
        a mini-batch

    Returns
    -------
    int: the number of samples in the mini-batch
    """
    if isinstance(batch, tuple):
        return _length_of_batch(batch[0])
    else:
        return len(batch)