diff --git a/pyproject.toml b/pyproject.toml index ba8acb64..5c3ad040 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,7 +18,7 @@ dependencies = [ "array-api-compat>=1.13", "array-api-extra>=0.10", "optree>=0.19", - "orng>=0.2.0", + "orng>=0.2.1", ] dynamic = [ "version" diff --git a/src/sequince/steps/mutation/metropolis.py b/src/sequince/steps/mutation/metropolis.py index eb50b451..28e1bc9a 100644 --- a/src/sequince/steps/mutation/metropolis.py +++ b/src/sequince/steps/mutation/metropolis.py @@ -6,7 +6,7 @@ import array_api_extra as xpx from ..._typing import Array -from ...array_api import asarray, namespace, to_numpy +from ...array_api import asarray, get_device, namespace, to_numpy from ...target import EvaluationCounts from ...transforms import ParticleTransform, select_transform_auxiliary from ..base import validate_n_steps @@ -168,7 +168,9 @@ def apply_with_mutation_state( ) empty_chain = ( xp.empty( - (0, *state.particles.shape), dtype=state.particles.dtype + (0, *state.particles.shape), + dtype=state.particles.dtype, + device=get_device(state.particles), ) if self._store_chains else None @@ -531,7 +533,9 @@ def apply_with_mutation_state( ) empty_chain = ( xp.empty( - (0, *state.particles.shape), dtype=state.particles.dtype + (0, *state.particles.shape), + dtype=state.particles.dtype, + device=get_device(state.particles), ) if self._store_chains else None @@ -588,7 +592,9 @@ def apply_with_mutation_state( acceptance_by_half: list[tuple[float, float]] = [] chain: list[Array] = [] for _ in range(n_steps): - accepted_population = xp.zeros((n_particles,), dtype=bool) + accepted_population = xp.zeros_like( + current_log_prob, dtype=xp.bool + ) group_a, group_b = self._split_indices( n_particles, rng, diff --git a/src/sequince/transforms.py b/src/sequince/transforms.py index 48a04d4b..af4109a6 100644 --- a/src/sequince/transforms.py +++ b/src/sequince/transforms.py @@ -331,8 +331,10 @@ def _log_det_entries(value, arr, xp): A vector of log-determinant entries, one per particle. """ if arr.ndim == 1 or arr.shape[0] == 1: - return xp.asarray(value) - return xp.full((arr.shape[0],), value, dtype=arr.dtype) + return asarray(value, xp=xp, like=arr) + # Broadcast on the particle device without extracting a host scalar. + value = asarray(value, xp=xp, dtype=arr.dtype, like=arr) + return xp.broadcast_to(value, (arr.shape[0],)) class ParticleTransform: