MAINT: move sequince._typing to sequince.typing

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
 
 

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,

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):

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

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

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

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(

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]]

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

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
 

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:

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
 

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:

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

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:

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

src/sequince/steps/mutation/blackjax.py

@@ -8,14 +8,14 @@
 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 (
     ParticleTransform,
     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

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

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 (

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 (

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
 

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
 

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

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
 

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

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
 
 

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

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

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

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
 

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__)