def _get_prediction_batch_shape(predictions):
    """ Get the batch size of the prediction.
    
    Args:
        predictions: a batch of predictions of any type. 
    
    Returns:
        int: the batch size of the prediction. 
    """
    if hasattr(predictions, 'batch_shape'):
        return predictions.batch_shape[0]
    if hasattr(predictions, 'shape'):   # If predictions has attribute shape 
        return predictions.shape[0]
    if issubclass(type(predictions), dict):
        return _get_prediction_batch_shape(predictions[next(iter(predictions))])
    assert False, "Cannot get the batch shape of prediction, unrecognized type"