<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def estimate_tau_exp(chains, **kwargs):
 """ Estimate the exponential auto-correlation time for all parameters in a chain. """


    # Calculate the normalised autocorrelation function in each parameter.
    rho = np.nan * np.ones(chains.shape[1:])
    for i in range(chains.shape[2]):
        try:
            rho[:, i] = autocorr.function(np.mean(chains[:, :, i], axis=0),
                **kwargs)
        except:
            continue

    # Take the max rho at any step.
    rho_max = np.max(rho, axis=1)

    # Now fit the max rho with an exponential profile.
    x = np.arange(rho_max.size)
    func = lambda tau_exp: np.exp(-x/tau_exp)
    chi = lambda tau_exp: func(tau_exp[0]) - rho_max # tau_exp is a list

    # Start with 50% of the chain length. probably OK.
    tau_exp, ier = leastsq(chi, [chains.shape[1]/2.])
    return (tau_exp, rho, func(tau_exp))