Calculating consecutive days in a 3D array

[chai0404]

I am trying to calculate the duration of consecutive days for a multidimensional array (latitude, longitude and time).
I have a piece of code that works for a single grid point (1D array) but would like for it to work in a 3D array.

def event_durations(event_mask, dim='time'):
    """
    Returns the lengths of events marked by 'event_mask'. Where 'event_mask' is 
    true an event is deemed active
    """
    event_stats = []
    
    assert event_mask.ndim == 1, "Only 1d arrays are implemented"

    # Loop over each entry, adding records for events to 'event_stats'
    current_event = None
    for i in range(event_mask.sizes[dim]):
        event_active = event_mask.isel({dim:i})
    
        if event_active:
            if not current_event:
                # A new event
                current_event = {'start': event_mask[dim].data[i], 'duration': 1}
            else:
                # An existing event
                current_event['duration'] += 1
        else:
            if current_event:
                # Event has finished, add to the record array
                event_stats.append(current_event)
                current_event = None

    if current_event:
        # Event active at the end
        event_stats.append(current_event)

    return pandas.DataFrame.from_records(event_stats, index='start')

I tried using the following function but that didn't work. How can I make the event_durations function work for a 3d xarray?

thw_out = xr.apply_ufunc(
    event_durations,          # the function name 
    thw_in,                             # the parameters of the function in order 
                                     # (only one, using defaults for others
    input_core_dims=[["time"]],      # thw should retain the 'time' axis
    output_core_dims=[['time']],     # the returned array's axis is the 'time' axis
    vectorize=True                   # all other axes should be looped over
)

Any guidance is appreciated.