|
64 | 64 | )
|
65 | 65 | from pymc.model import Model, modelcontext
|
66 | 66 | from pymc.parallel_sampling import Draw, _cpu_count
|
67 |
| -from pymc.step_methods import ( |
68 |
| - NUTS, |
69 |
| - BinaryGibbsMetropolis, |
70 |
| - BinaryMetropolis, |
71 |
| - CategoricalGibbsMetropolis, |
72 |
| - CompoundStep, |
73 |
| - DEMetropolis, |
74 |
| - HamiltonianMC, |
75 |
| - Metropolis, |
76 |
| - Slice, |
77 |
| -) |
| 67 | +from pymc.step_methods import NUTS, CompoundStep, DEMetropolis |
78 | 68 | from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
|
79 | 69 | from pymc.step_methods.hmc import quadpotential
|
80 | 70 | from pymc.util import (
|
|
98 | 88 | "draw",
|
99 | 89 | ]
|
100 | 90 |
|
101 |
| -STEP_METHODS = ( |
102 |
| - NUTS, |
103 |
| - HamiltonianMC, |
104 |
| - Metropolis, |
105 |
| - BinaryMetropolis, |
106 |
| - BinaryGibbsMetropolis, |
107 |
| - Slice, |
108 |
| - CategoricalGibbsMetropolis, |
109 |
| -) |
110 | 91 | Step: TypeAlias = Union[BlockedStep, CompoundStep]
|
111 | 92 |
|
112 | 93 | ArrayLike: TypeAlias = Union[np.ndarray, List[float]]
|
@@ -164,7 +145,7 @@ def instantiate_steppers(
|
164 | 145 | return steps
|
165 | 146 |
|
166 | 147 |
|
167 |
| -def assign_step_methods(model, step=None, methods=STEP_METHODS, step_kwargs=None): |
| 148 | +def assign_step_methods(model, step=None, methods=None, step_kwargs=None): |
168 | 149 | """Assign model variables to appropriate step methods.
|
169 | 150 |
|
170 | 151 | Passing a specified model will auto-assign its constituent stochastic
|
@@ -197,6 +178,9 @@ def assign_step_methods(model, step=None, methods=STEP_METHODS, step_kwargs=None
|
197 | 178 | steps = []
|
198 | 179 | assigned_vars = set()
|
199 | 180 |
|
| 181 | + if methods is None: |
| 182 | + methods = pm.STEP_METHODS |
| 183 | + |
200 | 184 | if step is not None:
|
201 | 185 | try:
|
202 | 186 | steps += list(step)
|
@@ -481,29 +465,7 @@ def sample(
|
481 | 465 | draws += tune
|
482 | 466 |
|
483 | 467 | initial_points = None
|
484 |
| - if step is None and init is not None and all_continuous(model.value_vars): |
485 |
| - try: |
486 |
| - # By default, try to use NUTS |
487 |
| - _log.info("Auto-assigning NUTS sampler...") |
488 |
| - initial_points, step = init_nuts( |
489 |
| - init=init, |
490 |
| - chains=chains, |
491 |
| - n_init=n_init, |
492 |
| - model=model, |
493 |
| - seeds=random_seed, |
494 |
| - progressbar=progressbar, |
495 |
| - jitter_max_retries=jitter_max_retries, |
496 |
| - tune=tune, |
497 |
| - initvals=initvals, |
498 |
| - **kwargs, |
499 |
| - ) |
500 |
| - except (AttributeError, NotImplementedError, tg.NullTypeGradError): |
501 |
| - # gradient computation failed |
502 |
| - _log.info("Initializing NUTS failed. Falling back to elementwise auto-assignment.") |
503 |
| - _log.debug("Exception in init nuts", exc_info=True) |
504 |
| - step = assign_step_methods(model, step, step_kwargs=kwargs) |
505 |
| - else: |
506 |
| - step = assign_step_methods(model, step, step_kwargs=kwargs) |
| 468 | + step = assign_step_methods(model, step, methods=pm.STEP_METHODS, step_kwargs=kwargs) |
507 | 469 |
|
508 | 470 | if isinstance(step, list):
|
509 | 471 | step = CompoundStep(step)
|
|
0 commit comments