def deaugment_state(state_array, aug_dim=0):
    """
    Compute deaugmented state arrays from
    Augmented NODE predictions
    Input::
    state_array: ANODE prediction state array
    aug_dim: How many rows to deaugment state vector
    Output::
    Deaugmented state array
    """
    if aug_dim == 0:
        return state_array
    else:
        return state_array[:,:-aug_dim]