<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def update_plot_limits(ax, white_space):
 """Sets the limit options of a matplotlib plot. Args: ax: matplotlib axes white_space(float):
 whitespace added to surround the tight limit of the data Note: This relies on ax.dataLim (in 2d) and ax.[xy, zz]_dataLim being set in 3d """


    if hasattr(ax, 'zz_dataLim'):
        bounds = ax.xy_dataLim.bounds
        ax.set_xlim(bounds[0] - white_space, bounds[0] + bounds[2] + white_space)
        ax.set_ylim(bounds[1] - white_space, bounds[1] + bounds[3] + white_space)

        bounds = ax.zz_dataLim.bounds
        ax.set_zlim(bounds[0] - white_space, bounds[0] + bounds[2] + white_space)
    else:
        bounds = ax.dataLim.bounds
        assert not any(map(np.isinf, bounds)), 'Cannot set bounds if dataLim has infinite elements'
        ax.set_xlim(bounds[0] - white_space, bounds[0] + bounds[2] + white_space)
        ax.set_ylim(bounds[1] - white_space, bounds[1] + bounds[3] + white_space)