Improve docs and add threadpoolctl to environments

Author: aseyboldt

conda-envs/environment-dev.yml

@@ -18,6 +18,7 @@ dependencies:
 - networkx
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
+- threadpoolctl>=3.1.0
 # Extra dependencies for dev, testing and docs build
 - ipython>=7.16
 - jax

conda-envs/environment-docs.yml

@@ -16,6 +16,7 @@ dependencies:
 - rich>=13.7.1
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
+- threadpoolctl>=3.1.0
 # Extra dependencies for docs build
 - ipython>=7.16
 - jax

conda-envs/environment-jax.yml

@@ -24,6 +24,7 @@ dependencies:
 - python-graphviz
 - networkx
 - rich>=13.7.1
+- threadpoolctl>=3.1.0
 # JAX is only compatible with Scipy 1.13.0 from >=0.4.26, but the respective version of
 # JAXlib is still not on conda: https://github.com/conda-forge/jaxlib-feedstock/pull/243
 - scipy>=1.4.1,<1.13.0

conda-envs/environment-test.yml

@@ -22,6 +22,7 @@ dependencies:
 - rich>=13.7.1
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
+- threadpoolctl>=3.1.0
 # Extra dependencies for testing
 - ipython>=7.16
 - pre-commit>=2.8.0

conda-envs/windows-environment-dev.yml

@@ -19,6 +19,7 @@ dependencies:
 - rich>=13.7.1
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
+- threadpoolctl>=3.1.0
 # Extra dependencies for dev, testing and docs build
 - ipython>=7.16
 - myst-nb

conda-envs/windows-environment-test.yml

@@ -22,6 +22,7 @@ dependencies:
 - rich>=13.7.1
 - scipy>=1.4.1
 - typing-extensions>=3.7.4
+- threadpoolctl>=3.1.0
 # Extra dependencies for testing
 - ipython>=7.16
 - pre-commit>=2.8.0

pymc/sampling/mcmc.py

@@ -505,12 +505,12 @@ def sample(
         This requires the chosen sampler to be installed.
         All samplers, except "pymc", require the full model to be continuous.
     blas_cores: int or "auto" or None, default = "auto"
-        The total number of threads blas and openmp functions should use during sampling. If set to None,
-        this will keep the default behavior of whatever blas implementation is used at runtime.
-        Setting it to "auto" will set it so that the total number of active blas threads is the
+        The total number of threads blas and openmp functions should use during sampling.
+        Setting it to "auto" will ensure that the total number of active blas threads is the
         same as the `cores` argument. If set to an integer, the sampler will try to use that total
         number of blas threads. If `blas_cores` is not divisible by `cores`, it might get rounded
-        down.
+        down. If set to None, this will keep the default behavior of whatever blas implementation
+        is used at runtime.
     initvals : optional, dict, array of dict
         Dict or list of dicts with initial value strategies to use instead of the defaults from
         `Model.initial_values`. The keys should be names of transformed random variables.
@@ -661,15 +661,6 @@ def sample(
 
     cores = min(cores, chains)
 
-    if cores < 1:
-        raise ValueError("`cores` must be larger or equal to one")
-
-    if chains < 1:
-        raise ValueError("`chains` must be larger or equal to one")
-
-    if blas_cores is not None and blas_cores < 1:
-        raise ValueError("`blas_cores` must be larger or equal to one")
-
     num_blas_cores_per_chain: int | None
     joined_blas_limiter: Callable[[], Any]