def METHOD_NAME(
    x: np.ndarray, shape: Sequence[int], A: np.ndarray, B: np.ndarray
) -> Tuple[float, np.ndarray]:
    """Compute the objective and gradient for NNLS"""
    # Scipy's lbfgs flattens all arrays, so we first reshape
    # the iterate x
    x = x.reshape(shape)
    # Compute the difference matrix
    diff = np.einsum("mf,...ft->...mt", A, x, optimize=True) - B
    # Compute the objective value
    value = (1 / B.size) * 0.5 * np.sum(diff**2)
    # And the gradient
    grad = (1 / B.size) * np.einsum("mf,...mt->...ft", A, diff, optimize=True)
    # Flatten the gradient
    return value, grad.flatten()