def METHOD_NAME(
    ctx, coord_scale, component_logits, component_scale, pis, which, shape, coeffs
):
    white = coord_scale.new(shape).normal_()
    which_component_scale = component_scale[which].unsqueeze(-1)
    z = coord_scale * which_component_scale * white
    ctx.save_for_backward(
        z, coord_scale, component_logits, component_scale, pis, coeffs
    )
    return z