Allow vector input for chain in generate_trace_plots