diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d2ac1ecf..02d3c8ab 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -106,7 +106,7 @@ when their types are not clear from initialization. - Use modern Python syntax supported by Python 3.11, such as `list[str]`, `dict[str, float]`, and `ParticleState | None`. Prefer collection interfaces such as `Mapping` or `Sequence` when callers need not supply a concrete container. -- Use the existing `Array` and `ArrayLike` aliases from `sequince._typing` for +- Use the existing `Array` and `ArrayLike` aliases from `sequince.typing` for backend-neutral numerical interfaces. Reserve backend-specific types such as `numpy.ndarray` for code that requires that backend. Use `ArrayLike` only when the function actually accepts and converts those inputs. @@ -136,7 +136,7 @@ meaning. Use `array` for backend-native arrays and `array_like` only when the function also accepts and converts inputs such as lists. ```python -from sequince._typing import Array +from sequince.typing import Array from sequince.array_api import namespace diff --git a/src/sequince/array_api.py b/src/sequince/array_api.py index 7ddef51f..80881b0e 100644 --- a/src/sequince/array_api.py +++ b/src/sequince/array_api.py @@ -21,7 +21,7 @@ ) from ._scipy_compat import enable_scipy_array_api as _enable_scipy_array_api -from ._typing import Array +from .typing import Array IS_NAMESPACE_FUNCTIONS = { "numpy": is_numpy_namespace, diff --git a/src/sequince/density.py b/src/sequince/density.py index b9e43e07..1a0e4e73 100644 --- a/src/sequince/density.py +++ b/src/sequince/density.py @@ -6,7 +6,6 @@ from numbers import Real from typing import Callable -from ._typing import Array from .steps.context import StepContext from .target import ( EvaluationCounts, @@ -17,6 +16,7 @@ TemperedTarget, _evaluate_component, ) +from .typing import Array class _DensityComponents(Target): diff --git a/src/sequince/filtering/kernel.py b/src/sequince/filtering/kernel.py index 4fe2acdb..383d076f 100644 --- a/src/sequince/filtering/kernel.py +++ b/src/sequince/filtering/kernel.py @@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Protocol -from .._typing import Array +from ..typing import Array if TYPE_CHECKING: from orng import RandomGenerator diff --git a/src/sequince/flows.py b/src/sequince/flows.py index d50816a9..79a8fae5 100644 --- a/src/sequince/flows.py +++ b/src/sequince/flows.py @@ -6,8 +6,8 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Protocol -from ._typing import Array from .array_api import asarray +from .typing import Array if TYPE_CHECKING: from .steps.context import StepContext diff --git a/src/sequince/mcmc.py b/src/sequince/mcmc.py index 7b28a1d2..15bb500d 100644 --- a/src/sequince/mcmc.py +++ b/src/sequince/mcmc.py @@ -4,7 +4,6 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Self -from ._typing import Array from .array_api import asarray, namespace from .checkpoint import CheckpointWriter, SMCCheckpoint from .result import SMCResult, _decode_result_output, _encode_result_output @@ -13,6 +12,7 @@ from .steps.base import StepContext from .steps.mutation import MutationStep, get_mutation_step from .transforms import ParticleTransform +from .typing import Array if TYPE_CHECKING: import numpy as np diff --git a/src/sequince/plot.py b/src/sequince/plot.py index 32853134..bc9763d5 100644 --- a/src/sequince/plot.py +++ b/src/sequince/plot.py @@ -8,8 +8,8 @@ import matplotlib.pyplot as plt import numpy as np -from ._typing import Array from .array_api import namespace, to_numpy +from .typing import Array def plot_chains( diff --git a/src/sequince/posterior.py b/src/sequince/posterior.py index 90d64c62..f3fc5041 100644 --- a/src/sequince/posterior.py +++ b/src/sequince/posterior.py @@ -4,7 +4,6 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Callable -from ._typing import Array from .proposals import ProposalDistribution from .state import ParticleState from .steps.context import LOG_PROPOSAL_OVERRIDE_KEY, StepContext @@ -16,15 +15,16 @@ _evaluate_component, ) from .transforms import ParticleTransform, get_prebuilt_transform +from .typing import Array from .utils.callables import call_with_optional_context from .utils.copying import array_safe_copy if TYPE_CHECKING: from orng import RandomGenerator - from ._typing import MutationStrategy from .result import SMCResult from .steps.context import StepContext + from .typing import MutationStrategy LogProposalLike = TargetComponentFn | ProposalDistribution JointLogDensityFn = Callable[..., tuple[Array | float, Array | float]] diff --git a/src/sequince/prebuilt_samplers.py b/src/sequince/prebuilt_samplers.py index 482a0fc7..955dfdaa 100644 --- a/src/sequince/prebuilt_samplers.py +++ b/src/sequince/prebuilt_samplers.py @@ -9,7 +9,6 @@ import numpy as np -from ._typing import MetaData, MutationStrategy from .checkpoint import CheckpointWriter from .posterior import TemperedPosterior from .proposals import ProposalDistribution @@ -42,6 +41,7 @@ from .steps.resampling import ResamplingStep from .steps.weighting import TargetESSRatio from .transforms import ParticleTransform, get_prebuilt_transform +from .typing import MetaData, MutationStrategy if TYPE_CHECKING: from .state import ParticleState diff --git a/src/sequince/proposals.py b/src/sequince/proposals.py index 43c2f641..286498dc 100644 --- a/src/sequince/proposals.py +++ b/src/sequince/proposals.py @@ -13,7 +13,6 @@ from array_api_compat import device as get_device -from ._typing import Array from .array_api import asarray, namespace, to_numpy from .flows import ( NormalizingFlow, @@ -22,6 +21,7 @@ flow_state_dict, load_flow_state_dict, ) +from .typing import Array from .utils.callables import call_with_optional_context from .utils.random import resolve_rng diff --git a/src/sequince/result.py b/src/sequince/result.py index 597d8e9d..82d68bb5 100644 --- a/src/sequince/result.py +++ b/src/sequince/result.py @@ -10,12 +10,12 @@ import numpy as np from numpy.random import Generator as NumPyGenerator -from ._typing import Array from .array_api import asarray, get_device, namespace, to_numpy from .state import ParticleState from .steps.constraints import _STOP_REASONS, NestedSamplingLevel from .steps.context import StepContext from .steps.meta import META_SMC_OUTPUT_KEY +from .typing import Array from .utils.random import resolve_rng if TYPE_CHECKING: diff --git a/src/sequince/state.py b/src/sequince/state.py index c81f2321..2a8cb054 100644 --- a/src/sequince/state.py +++ b/src/sequince/state.py @@ -9,10 +9,10 @@ from array_api_compat import device as get_device -from ._typing import Array, ArrayLike from .array_api import asarray, logsumexp, namespace from .array_api import copy as array_copy from .exceptions import DegenerateWeightsError +from .typing import Array, ArrayLike from .utils.copying import array_safe_copy from .utils.pytrees import pytree_dataclass diff --git a/src/sequince/stats.py b/src/sequince/stats.py index 841cb8ec..aa3d1f20 100644 --- a/src/sequince/stats.py +++ b/src/sequince/stats.py @@ -1,8 +1,8 @@ from __future__ import annotations -from ._typing import Array from .array_api import asarray, logsumexp, namespace from .state import normalize_log_weights +from .typing import Array def ess_from_log_weights(log_weights: Array) -> Array: diff --git a/src/sequince/steps/constraints.py b/src/sequince/steps/constraints.py index 93ff0183..c973c58e 100644 --- a/src/sequince/steps/constraints.py +++ b/src/sequince/steps/constraints.py @@ -10,7 +10,6 @@ import numpy as np from orng import RandomGenerator -from .._typing import Array from ..array_api import asarray, is_jax_namespace, logsumexp, namespace from ..array_api import copy as array_copy from ..state import ( @@ -24,6 +23,7 @@ TargetEvaluation, _evaluate_component, ) +from ..typing import Array from .base import SMCStep from .context import StepContext from .control import apply_functional_if diff --git a/src/sequince/steps/control.py b/src/sequince/steps/control.py index 222b6aa2..3cf4b427 100644 --- a/src/sequince/steps/control.py +++ b/src/sequince/steps/control.py @@ -7,8 +7,8 @@ from acai import cond -from .._typing import Array from ..array_api import is_jax_namespace, namespace +from ..typing import Array from .base import SMCStep, StepContext if TYPE_CHECKING: diff --git a/src/sequince/steps/mutation/base.py b/src/sequince/steps/mutation/base.py index 225662a9..368c1357 100644 --- a/src/sequince/steps/mutation/base.py +++ b/src/sequince/steps/mutation/base.py @@ -5,10 +5,10 @@ import numpy as np -from ..._typing import Array from ...array_api import asarray, namespace, to_numpy from ...target import EvaluationCounts, Target from ...transforms import ParticleTransform, _repeat_transform_auxiliary +from ...typing import Array from ...utils.callables import call_with_optional_context from ...utils.copying import array_safe_copy from ..base import SMCStep diff --git a/src/sequince/steps/mutation/blackjax.py b/src/sequince/steps/mutation/blackjax.py index b480757c..294e0a6e 100644 --- a/src/sequince/steps/mutation/blackjax.py +++ b/src/sequince/steps/mutation/blackjax.py @@ -8,7 +8,6 @@ import numpy as np from array_api_compat import is_jax_namespace -from ..._typing import Array from ...array_api import asarray, namespace, to_numpy from ...target import EvaluationCounts from ...transforms import ( @@ -16,6 +15,7 @@ has_particle_aligned_auxiliary, select_transform_auxiliary, ) +from ...typing import Array from ...utils.pytrees import pytree_dataclass from ..base import validate_n_steps from ..context import StepContext diff --git a/src/sequince/steps/mutation/blocked.py b/src/sequince/steps/mutation/blocked.py index 7fe38e82..1328220c 100644 --- a/src/sequince/steps/mutation/blocked.py +++ b/src/sequince/steps/mutation/blocked.py @@ -7,8 +7,8 @@ import numpy as np -from ..._typing import Array from ...transforms import ParticleTransform +from ...typing import Array from ..base import SMCStep from ..context import StepContext from .base import TransformRefit diff --git a/src/sequince/steps/mutation/emcee.py b/src/sequince/steps/mutation/emcee.py index e927f610..e817844e 100644 --- a/src/sequince/steps/mutation/emcee.py +++ b/src/sequince/steps/mutation/emcee.py @@ -8,13 +8,13 @@ import numpy as np -from ..._typing import Array from ...array_api import asarray, namespace, to_numpy from ...target import EvaluationCounts, TargetEvaluation from ...transforms import ( ParticleTransform, has_particle_aligned_auxiliary, ) +from ...typing import Array from ..base import validate_n_steps from ..context import StepContext from .base import ( diff --git a/src/sequince/steps/mutation/eryn.py b/src/sequince/steps/mutation/eryn.py index f1dce8ec..5871871e 100644 --- a/src/sequince/steps/mutation/eryn.py +++ b/src/sequince/steps/mutation/eryn.py @@ -8,13 +8,13 @@ import numpy as np from array_api_compat import is_cupy_namespace -from ..._typing import Array from ...array_api import asarray, namespace, to_numpy from ...target import EvaluationCounts from ...transforms import ( ParticleTransform, has_particle_aligned_auxiliary, ) +from ...typing import Array from ..base import validate_n_steps from ..context import StepContext from .base import ( diff --git a/src/sequince/steps/mutation/kernels.py b/src/sequince/steps/mutation/kernels.py index 3e7a7a7e..d3c228c2 100644 --- a/src/sequince/steps/mutation/kernels.py +++ b/src/sequince/steps/mutation/kernels.py @@ -9,9 +9,9 @@ import numpy as np from array_api_compat import device as get_device -from ..._typing import Array from ...array_api import asarray, namespace, to_numpy from ...target import EvaluationCounts +from ...typing import Array from ...utils.pytrees import pytree_dataclass from ..base import StepContext diff --git a/src/sequince/steps/mutation/metrics/base.py b/src/sequince/steps/mutation/metrics/base.py index 781d6881..bdcd49ab 100644 --- a/src/sequince/steps/mutation/metrics/base.py +++ b/src/sequince/steps/mutation/metrics/base.py @@ -6,8 +6,8 @@ from dataclasses import dataclass, replace from typing import Any, Protocol -from ...._typing import Array from ....array_api import asarray, namespace +from ....typing import Array from ...context import StepContext from .state import MutationMetricsOutput, MutationStepState, MutationTransition diff --git a/src/sequince/steps/mutation/metrics/de.py b/src/sequince/steps/mutation/metrics/de.py index d1be2681..eebf7dd7 100644 --- a/src/sequince/steps/mutation/metrics/de.py +++ b/src/sequince/steps/mutation/metrics/de.py @@ -5,8 +5,8 @@ from dataclasses import dataclass from typing import ClassVar -from ...._typing import Array from ....array_api import asarray, get_device, namespace +from ....typing import Array from ...context import StepContext from .esjd import _squared_jump_distances from .state import MutationTransition diff --git a/src/sequince/steps/mutation/metrics/esjd.py b/src/sequince/steps/mutation/metrics/esjd.py index 355b72eb..777b77d2 100644 --- a/src/sequince/steps/mutation/metrics/esjd.py +++ b/src/sequince/steps/mutation/metrics/esjd.py @@ -4,8 +4,8 @@ from dataclasses import dataclass -from ...._typing import Array from ....array_api import asarray, namespace +from ....typing import Array from ...context import StepContext from .state import MutationTransition diff --git a/src/sequince/steps/mutation/metrics/rjd.py b/src/sequince/steps/mutation/metrics/rjd.py index add5d323..b5fe283e 100644 --- a/src/sequince/steps/mutation/metrics/rjd.py +++ b/src/sequince/steps/mutation/metrics/rjd.py @@ -8,8 +8,8 @@ from orng import RandomGenerator -from ...._typing import Array from ....array_api import asarray, get_device, namespace, to_numpy +from ....typing import Array from ....utils.pytrees import pytree_dataclass from ...base import validate_n_steps from ...context import StepContext diff --git a/src/sequince/steps/mutation/metrics/state.py b/src/sequince/steps/mutation/metrics/state.py index 74ba11d6..265f619c 100644 --- a/src/sequince/steps/mutation/metrics/state.py +++ b/src/sequince/steps/mutation/metrics/state.py @@ -6,8 +6,8 @@ from dataclasses import dataclass, field from typing import Any -from ...._typing import Array from ....array_api import asarray, namespace +from ....typing import Array from ....utils.pytrees import pytree_dataclass diff --git a/src/sequince/steps/mutation/metropolis.py b/src/sequince/steps/mutation/metropolis.py index 6ea63383..89f7d91e 100644 --- a/src/sequince/steps/mutation/metropolis.py +++ b/src/sequince/steps/mutation/metropolis.py @@ -8,10 +8,10 @@ import array_api_extra as xpx import optree -from ..._typing import Array from ...array_api import asarray, get_device, namespace, to_numpy from ...target import EvaluationCounts from ...transforms import ParticleTransform, select_transform_auxiliary +from ...typing import Array from ...utils.pytrees import OPTREE_NAMESPACE from ..base import validate_n_steps from ..context import StepContext diff --git a/src/sequince/steps/mutation/minipcn.py b/src/sequince/steps/mutation/minipcn.py index 9aac59c6..c04752f8 100644 --- a/src/sequince/steps/mutation/minipcn.py +++ b/src/sequince/steps/mutation/minipcn.py @@ -5,10 +5,10 @@ from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any -from ..._typing import Array from ...array_api import asarray, namespace from ...target import EvaluationCounts from ...transforms import ParticleTransform +from ...typing import Array from ...utils.pytrees import pytree_dataclass from ..base import validate_n_steps from ..context import StepContext diff --git a/src/sequince/steps/persistent.py b/src/sequince/steps/persistent.py index c33dc00f..ec2a3e64 100644 --- a/src/sequince/steps/persistent.py +++ b/src/sequince/steps/persistent.py @@ -10,7 +10,6 @@ from numbers import Integral from typing import TYPE_CHECKING, Any -from .._typing import Array, ArrayLike from ..array_api import asarray, logsumexp, namespace from ..state import ( ParticleState, @@ -19,6 +18,7 @@ ) from ..stats import ess_from_log_weights from ..target import EvaluationCounts, Target, TemperedTarget +from ..typing import Array, ArrayLike from ..utils.optimize import bisect_monotone_decreasing from .base import SMCStep from .context import StepContext diff --git a/src/sequince/steps/resampling.py b/src/sequince/steps/resampling.py index bfae066e..2692186a 100644 --- a/src/sequince/steps/resampling.py +++ b/src/sequince/steps/resampling.py @@ -9,9 +9,9 @@ import array_api_extra as xpx from array_api_compat import device as get_device -from .._typing import Array from ..array_api import asarray, namespace from ..state import normalize_log_weights +from ..typing import Array from ..utils.pytrees import pytree_dataclass from .base import SMCStep, StepContext diff --git a/src/sequince/steps/transformations.py b/src/sequince/steps/transformations.py index 7bf84d1b..eab63eaf 100644 --- a/src/sequince/steps/transformations.py +++ b/src/sequince/steps/transformations.py @@ -15,9 +15,9 @@ if TYPE_CHECKING: from orng import RandomGenerator - from .._typing import Array from ..posterior import TemperedPosterior from ..state import ParticleState + from ..typing import Array logger = logging.getLogger(__name__) diff --git a/src/sequince/steps/weighting.py b/src/sequince/steps/weighting.py index a7b5b402..a6c31ab8 100644 --- a/src/sequince/steps/weighting.py +++ b/src/sequince/steps/weighting.py @@ -6,7 +6,6 @@ from dataclasses import replace from typing import TYPE_CHECKING, Any, Callable, Literal, cast -from .._typing import Array from ..array_api import asarray, logsumexp, namespace, to_numpy from ..state import _validate_log_normalizer, normalize_log_weights from ..stats import ( @@ -19,6 +18,7 @@ SequentialTarget, TemperedTarget, ) +from ..typing import Array from ..utils.callables import call_with_optional_context from ..utils.optimize import bisect_monotone_decreasing from .base import SMCStep, StepContext diff --git a/src/sequince/target.py b/src/sequince/target.py index 123d6cfc..7245db23 100644 --- a/src/sequince/target.py +++ b/src/sequince/target.py @@ -11,9 +11,9 @@ import array_api_extra as xpx import numpy as np -from ._typing import Array from .array_api import asarray, namespace from .array_api import copy as array_copy +from .typing import Array from .utils.callables import call_with_optional_context from .utils.pytrees import pytree_dataclass diff --git a/src/sequince/transforms.py b/src/sequince/transforms.py index c41f5e1b..a3ec2866 100644 --- a/src/sequince/transforms.py +++ b/src/sequince/transforms.py @@ -14,7 +14,6 @@ from scipy.special import ndtr, ndtri from scipy.stats import circmean -from ._typing import Array from .array_api import asarray, namespace, to_numpy from .flows import ( NormalizingFlow, @@ -23,6 +22,7 @@ flow_state_dict, load_flow_state_dict, ) +from .typing import Array from .utils.pytrees import OPTREE_NAMESPACE, pytree_dataclass from .utils.random import RNGLike, resolve_rng diff --git a/src/sequince/_typing.py b/src/sequince/typing.py similarity index 77% rename from src/sequince/_typing.py rename to src/sequince/typing.py index 4617b294..5c533d04 100644 --- a/src/sequince/_typing.py +++ b/src/sequince/typing.py @@ -1,3 +1,9 @@ +"""Type aliases for backend-neutral numerical interfaces. + +``Array`` and ``ArrayLike`` are the supported public aliases. Array shapes, +dtype constraints, and accepted input conversions are specified by each API. +""" + from __future__ import annotations from collections.abc import Sequence diff --git a/src/sequince/utils/bounds.py b/src/sequince/utils/bounds.py index 4b455ab2..13e70d0f 100644 --- a/src/sequince/utils/bounds.py +++ b/src/sequince/utils/bounds.py @@ -2,8 +2,8 @@ from typing import Any -from .._typing import Array from ..array_api import namespace +from ..typing import Array def in_bounds(particles: Array, bounds: Array, xp: Any | None = None) -> Array: diff --git a/src/sequince/utils/optimize.py b/src/sequince/utils/optimize.py index 2a72bdeb..0e73ee95 100644 --- a/src/sequince/utils/optimize.py +++ b/src/sequince/utils/optimize.py @@ -5,8 +5,8 @@ from acai import while_loop from array_api_compat import is_jax_namespace -from .._typing import Array from ..array_api import namespace +from ..typing import Array def bisect_monotone_decreasing( diff --git a/tests/test_functional_smc.py b/tests/test_functional_smc.py index 136c6b83..f316f770 100644 --- a/tests/test_functional_smc.py +++ b/tests/test_functional_smc.py @@ -26,7 +26,7 @@ @pytest.mark.parametrize("n_particles", [16, 17]) def test_compiled_differential_evolution_replays(n_particles: int) -> None: from sequince import make_standard_bayesian_smc - from sequince._typing import Array + from sequince.typing import Array def prior(particles: Array, context: StepContext) -> Array: return -0.5 * jnp.sum(particles**2, axis=-1) - jnp.log(2 * jnp.pi)