Optimize the following code to use the variable: variance in the code. def e_step(xs: np.ndarray, ys: np.ndarray, affine: np.ndarray, translation: np.ndarray, variance: float) -> np.ndarray: """ The e-step of the em algorithm, estimating the responsibility P=[p(y_m | x_n)] based on current model :param xs: a set of points with size (N, D), N is the number of samples, D is the dimension of points :param ys: a set of points with size (M, D), M is the number of samples, D is the dimension of points :param affine: an affine matrix with size (D, D) :param translation: a translation vector with size (1, D) :param variance: a float controlling the variance of each Gaussian component :return: the responsibility matrix P=[p(y_m | x_n)] with size (N, M), which row is the conditional probability of clusters given the n-th sample x_n """ # TODO: Change the code below and implement the E-step of GMM responsibility = np.ones((xs.shape[0], ys.shape[0])) / ys.shape[0] for n in range(xs.shape[0]): for m in range(ys.shape[0]): temp = -0.5 * np.linalg.norm(xs[n] - ys[m] @ affine - translation) ** 2 responsibility[n, m] = 1 / (2 * np.pi) ** (xs.shape[1] / 2) * np.exp(temp) return responsibility / np.sum(responsibility, axis=1, keepdims=True)

Here's an optimized version of the code that uses the `variance` variable:

def e_step(xs: np.ndarray, ys: np.ndarray, affine: np.ndarray, translation: np.ndarray, variance: float) -> np.ndarray:
    """
    The e-step of the em algorithm, estimating the responsibility P=[p(y_m | x_n)] based on current model
    :param xs: a set of points with size (N, D), N is the number of samples, D is the dimension of points
    :param ys: a set of points with size (M, D), M is the number of samples, D is the dimension of points
    :param affine: an affine matrix with size (D, D)
    :param translation: a translation vector with size (1, D)
    :param variance: a float controlling the variance of each Gaussian component
    :return:
        the responsibility matrix P=[p(y_m | x_n)] with size (N, M),
        which row is the conditional probability of clusters given the n-th sample x_n
    """
    # Calculate the constant term in the Gaussian distribution
    constant_term = 1 / (2 * np.pi * variance) ** (xs.shape[1] / 2)
    
    # Calculate the Mahalanobis distance between xs and ys
    mahalanobis = np.sum((xs[:, np.newaxis] - ys) @ affine.T ** 2, axis=2)
    
    # Calculate the responsibility matrix
    responsibility = constant_term * np.exp(-mahalanobis / (2 * variance))
    
    # Normalize the responsibility matrix
    return responsibility / np.sum(responsibility, axis=1, keepdims=True)

In this optimized version, we pre-calculate the constant term in the Gaussian distribution and the Mahalanobis distance between `xs` and `ys`. Then, we use these pre-calculated values to calculate the responsibility matrix in a vectorized way, which is much faster than the nested loops in the original code. Finally, we normalize the responsibility matrix and return it.