TF_ASSIGN_OR_RETURN(const XlaFunctionInfo xla_function_info,
                        get_xla_function_info(*launch));

    std::vector<const Edge*> in_edges;
    TF_RETURN_IF_ERROR(launch->input_edges(&in_edges));

    const int num_inputs = in_edges.size();
    const int variable_start_index = xla_function_info.variable_start_index;
    const int num_variables = num_inputs - variable_start_index;