<commit_before>from . import backends
from .backends.base import merge_traces, BaseTrace, MultiTrace
from .backends.ndarray import NDArray
import multiprocessing as mp
from time import time
from .core import *
from .step_methods import *
from .progressbar import progress_bar
from numpy.random import seed

__all__ = ['sample', 'iter_sample']

def assign_step_methods(model, step,
        methods=(NUTS, HamiltonianMC, Metropolis, BinaryMetropolis, 
        Slice, ElemwiseCategoricalStep)):
    
    steps = []
    assigned_vars = set()
    if step is not None:
        steps = np.append(steps, step).tolist()
        try:
            assigned_vars = assigned_vars | set(step.vars)
        except AttributeError:
            for s in step.methods:
                assigned_vars = assigned_vars | set(s.vars)
    
    # Use competence classmethods to select step methods for remaining variables
    selected_steps = {s:[] for s in methods}
    for var in model.free_RVs:
        if not var in assigned_vars:
               
            competences = {s:s._competence(var) for s in methods}
            
            selected = max(competences.keys(), key=(lambda k: competences[k]))
        
            print('Assigned {0} to {1}'.format(selected, var))
            selected_steps[selected].append(var)
    
    # Instantiate all selected step methods
    steps += [s(vars=selected_steps[s]) for s in selected_steps if selected_steps[s]]
    
    if len(steps)==1:
        steps = steps[0]    
                <commit_msg>Fixed step assignment failure in sample()<commit_after>from . import backends
from .backends.base import merge_traces, BaseTrace, MultiTrace
from .backends.ndarray import NDArray
import multiprocessing as mp
from time import time
from .core import *
from .step_methods import *
from .progressbar import progress_bar
from numpy.random import seed

__all__ = ['sample', 'iter_sample']

def assign_step_methods(model, step,
        methods=(NUTS, HamiltonianMC, Metropolis, BinaryMetropolis, 
        Slice, ElemwiseCategoricalStep)):
        
    steps = []
    assigned_vars = set()
    if step is not None:
        steps = np.append(steps, step).tolist()
        for s in steps:
            try:
                assigned_vars = assigned_vars | set(s.vars)
            except AttributeError:
                for m in s.methods:
                    assigned_vars = assigned_vars | set(m.vars)
    
    # Use competence classmethods to select step methods for remaining variables
    selected_steps = {s:[] for s in methods}
    for var in model.free_RVs:
        if not var in assigned_vars:
               
            competences = {s:s._competence(var) for s in methods}
            
            selected = max(competences.keys(), key=(lambda k: competences[k]))
        
            print('Assigned {0} to {1}'.format(selected, var))
            selected_steps[selected].append(var)
    
    # Instantiate all selected step methods
    steps += [s(vars=selected_steps[s]) for s in selected_steps if selected_steps[s]]
    
    if len(steps)==1:
        steps = steps[0]    
                
{'fraction_non_alphanumeric': 0.08638407193320488, 'fraction_numerical': 0.002569043031470777, 'mean_word_length': 2.4611111111111112, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3577534', 'n_tokens_mistral': 887, 'n_tokens_neox': 875, 'n_words': 269}