Merge pull request #340 from jakevdp/row_stack

Replace jnp.row_stack with jnp.vstack

dynamax/generalized_gaussian_ssm/inference.py

@@ -342,8 +342,8 @@ def _step(carry, args):
     _, (smoothed_means, smoothed_covs) = lax.scan(_step, init_carry, args)
 
     # Reverse the arrays and return
-    smoothed_means = jnp.row_stack((smoothed_means[::-1], filtered_means[-1][None, ...]))
-    smoothed_covs = jnp.row_stack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
+    smoothed_means = jnp.vstack((smoothed_means[::-1], filtered_means[-1][None, ...]))
+    smoothed_covs = jnp.vstack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
     return PosteriorGSSMSmoothed(
         marginal_loglik=ll,
         filtered_means=filtered_means,

dynamax/hidden_markov_model/inference.py

@@ -303,7 +303,7 @@ def _step(carry, args):
     _, rev_smoothed_probs = lax.scan(_step, carry, args)
 
     # Reverse the arrays and return
-    smoothed_probs = jnp.row_stack([rev_smoothed_probs[::-1], filtered_probs[-1]])
+    smoothed_probs = jnp.vstack([rev_smoothed_probs[::-1], filtered_probs[-1]])
 
     # Package into a posterior
     posterior = HMMPosterior(

dynamax/hidden_markov_model/inference_test.py

@@ -79,7 +79,7 @@ def test_hmm_filter(key=0, num_timesteps=3, num_states=2):
     # Compare predicted_probs to manually computed entries
     for t in range(num_timesteps):
         log_joint_t = big_log_joint(initial_probs, transition_matrix,
-                                    jnp.row_stack([log_lkhds[:t], jnp.zeros(num_states)]))
+                                    jnp.vstack([log_lkhds[:t], jnp.zeros(num_states)]))
 
         log_joint_t -= logsumexp(log_joint_t)
         predicted_probs_t = jnp.exp(logsumexp(log_joint_t, axis=tuple(jnp.arange(t))))

dynamax/hidden_markov_model/models/arhmm.py

@@ -187,14 +187,14 @@ def _step(carry, key):
             key1, key2 = jr.split(key, 2)
             state = self.transition_distribution(params, prev_state).sample(seed=key2)
             emission = self.emission_distribution(params, state, inputs=jnp.ravel(prev_emissions)).sample(seed=key1)
-            next_prev_emissions = jnp.row_stack([emission, prev_emissions[:-1]])
+            next_prev_emissions = jnp.vstack([emission, prev_emissions[:-1]])
             return (state, next_prev_emissions), (state, emission)
 
         # Sample the initial state
         key1, key2, key = jr.split(key, 3)
         initial_state = self.initial_distribution(params).sample(seed=key1)
         initial_emission = self.emission_distribution(params, initial_state, inputs=jnp.ravel(prev_emissions)).sample(seed=key2)
-        initial_prev_emissions = jnp.row_stack([initial_emission, prev_emissions[:-1]])
+        initial_prev_emissions = jnp.vstack([initial_emission, prev_emissions[:-1]])
 
         # Sample the remaining emissions and states
         next_keys = jr.split(key, num_timesteps - 1)

dynamax/linear_gaussian_ssm/inference.py

@@ -544,8 +544,8 @@ def _step(carry, args):
     _, (smoothed_means, smoothed_covs, smoothed_cross) = lax.scan(_step, init_carry, args)
 
     # Reverse the arrays and return
-    smoothed_means = jnp.row_stack((smoothed_means[::-1], filtered_means[-1][None, ...]))
-    smoothed_covs = jnp.row_stack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
+    smoothed_means = jnp.vstack((smoothed_means[::-1], filtered_means[-1][None, ...]))
+    smoothed_covs = jnp.vstack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
     smoothed_cross = smoothed_cross[::-1]
     return PosteriorGSSMSmoothed(
         marginal_loglik=ll,
@@ -610,5 +610,5 @@ def _step(carry, args):
         jnp.arange(num_timesteps - 2, -1, -1),
     )
     _, reversed_states = lax.scan(_step, last_state, args)
-    states = jnp.row_stack([reversed_states[::-1], last_state])
+    states = jnp.vstack([reversed_states[::-1], last_state])
     return states

dynamax/linear_gaussian_ssm/info_inference.py

@@ -276,8 +276,8 @@ def _smooth_step(carry, args):
     _, (smoothed_etas, smoothed_precisions) = lax.scan(_smooth_step, init_carry, args)
 
     # Reverse the arrays and return
-    smoothed_etas = jnp.row_stack((smoothed_etas[::-1], filtered_etas[-1][None, ...]))
-    smoothed_precisions = jnp.row_stack((smoothed_precisions[::-1], filtered_precisions[-1][None, ...]))
+    smoothed_etas = jnp.vstack((smoothed_etas[::-1], filtered_etas[-1][None, ...]))
+    smoothed_precisions = jnp.vstack((smoothed_precisions[::-1], filtered_precisions[-1][None, ...]))
     return PosteriorGSSMInfoSmoothed(
         marginal_loglik=ll,
         filtered_etas=filtered_etas,

dynamax/nonlinear_gaussian_ssm/inference_ekf.py

@@ -242,8 +242,8 @@ def _step(carry, args):
     _, (smoothed_means, smoothed_covs) = lax.scan(_step, init_carry, args)
 
     # Reverse the arrays and return
-    smoothed_means = jnp.row_stack((smoothed_means[::-1], filtered_means[-1][None, ...]))
-    smoothed_covs = jnp.row_stack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
+    smoothed_means = jnp.vstack((smoothed_means[::-1], filtered_means[-1][None, ...]))
+    smoothed_covs = jnp.vstack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
     return PosteriorGSSMSmoothed(
         marginal_loglik=ll,
         filtered_means=filtered_means,
@@ -311,7 +311,7 @@ def _step(carry, args):
         jnp.arange(num_timesteps - 2, -1, -1),
     )
     _, reversed_states = lax.scan(_step, last_state, args)
-    states = jnp.row_stack([reversed_states[::-1], last_state])
+    states = jnp.vstack([reversed_states[::-1], last_state])
     return states
 
 

dynamax/nonlinear_gaussian_ssm/inference_ukf.py

@@ -276,8 +276,8 @@ def _step(carry, args):
     _, (smoothed_means, smoothed_covs) = lax.scan(_step, init_carry, args)
 
     # Reverse the arrays and return
-    smoothed_means = jnp.row_stack((smoothed_means[::-1], filtered_means[-1][None, ...]))
-    smoothed_covs = jnp.row_stack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
+    smoothed_means = jnp.vstack((smoothed_means[::-1], filtered_means[-1][None, ...]))
+    smoothed_covs = jnp.vstack((smoothed_covs[::-1], filtered_covs[-1][None, ...]))
     return PosteriorGSSMSmoothed(
         marginal_loglik=ll,
         filtered_means=filtered_means,