GaussianHMM: TypeError: reshape() got an unexpected keyword argument 'newshape'

[rmndrs89]

Dynamax: 1.0.1
JAX: 0.8.2

I am simulating some time series data that are generated from two hypothetical states (num_states=2). The time series measures whether an electronic appliance is off or on (0V or 5V, respectively), simulated with noise with a small amplitude:

fs = 64. # sampling frequency (Hz)
num_timesteps = 3_840
num_states = 2
ts = np.arange(num_timesteps) / fs  # time (s)

zs = np.zeros_like(ts)
ys = np.random.normal(loc=0.0, scale=0.05, size=num_timesteps)

current_state = 0
transition_probs = np.array([[0.2, 0.8], [0.3, 0.7]])
for t in range(0, num_timesteps - 64 + 1, 64):
    zs[t:t+64] = current_state
    current_state = np.random.choice(num_states, replace=True, p=transition_probs[current_state])
ys += 5 * zs

Then I setup a Gaussian Hidden Markov Model:

import jax.random as jr
from dynamax.hidden_markov_model import GaussianHMM
hmm = GaussianHMM(num_states=num_states, emission_dim=1)
hmm_params, hmm_props = hmm.initialize(key=jr.PRNGKey(0), method="kmeans", emissions=ys)
# hmm_params, hmm_props = hmm.initialize(key=jr.PRNGKey(0), method="kmeans")

but I get the following error, that I don't fully understand. Do my inputs have the wrong shape? Or is the initialization not correct?

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[6], [line 1](vscode-notebook-cell:?execution_count=6&line=1)
----> [1](vscode-notebook-cell:?execution_count=6&line=1) true_params, _ = hmm.initialize(
      2     initial_probs=initial_probs,
      3     transition_matrix=transition_matrix,
      4     emission_means=emission_means,
      5     emission_covariances=emissions_covs
      6 )

File ~/Projects/pytsa/.venv/lib/python3.12/site-packages/dynamax/hidden_markov_model/models/gaussian_hmm.py:808, in GaussianHMM.initialize(self, key, method, initial_probs, transition_matrix, emission_means, emission_covariances, emissions)
    806 key1, key2, key3 = jr.split(key , 3)
    807 params, props = dict(), dict()
--> [808](https://file+.vscode-resource.vscode-cdn.net/home/robbin/Projects/pytsa/notebooks/~/Projects/pytsa/.venv/lib/python3.12/site-packages/dynamax/hidden_markov_model/models/gaussian_hmm.py:808) params["initial"], props["initial"] = self.initial_component.initialize(
    809     key1, method=method, initial_probs=initial_probs
    810 )
    811 params["transitions"], props["transitions"] = self.transition_component.initialize(
    812     key2, method=method, transition_matrix=transition_matrix
    813 )
    814 params["emissions"], props["emissions"] = self.emission_component.initialize(
    815     key3, method=method, emission_means=emission_means, emission_covariances=emission_covariances, emissions=emissions
    816 )

File ~/Projects/pytsa/.venv/lib/python3.12/site-packages/dynamax/hidden_markov_model/models/initial.py:64, in StandardHMMInitialState.initialize(self, key, method, initial_probs)
     62 # Package the results into dictionaries
     63 params = ParamsStandardHMMInitialState(probs=initial_probs)
---> [64](https://file+.vscode-resource.vscode-cdn.net/home/robbin/Projects/pytsa/notebooks/~/Projects/pytsa/.venv/lib/python3.12/site-packages/dynamax/hidden_markov_model/models/initial.py:64) props = ParamsStandardHMMInitialState(probs=ParameterProperties(constrainer=tfb.SoftmaxCentered()))
     65 return params, props

File ~/Projects/pytsa/.venv/lib/python3.12/site-packages/tensorflow_probability/substrates/jax/bijectors/softmax_centered.py:65, in SoftmaxCentered.__init__(self, validate_args, name)
     63 parameters = dict(locals())
     64 with tf.name_scope(name) as name:
---> [65](https://file+.vscode-resource.vscode-cdn.net/home/robbin/Projects/pytsa/notebooks/~/Projects/pytsa/.venv/lib/python3.12/site-packages/tensorflow_probability/substrates/jax/bijectors/softmax_centered.py:65)   self._pad = pad_lib.Pad(validate_args=validate_args)
     66   super(SoftmaxCentered, self).__init__(
     67       forward_min_event_ndims=1,
     68       validate_args=validate_args,
     69       parameters=parameters,
     70       name=name)

File ~/Projects/pytsa/.venv/lib/python3.12/site-packages/tensorflow_probability/substrates/jax/bijectors/pad.py:143, in Pad.__init__(self, paddings, mode, constant_values, axis, validate_args, name)
    139 self._mode = mode
    140 self._constant_values = tensor_util.convert_nonref_to_tensor(
    141     constant_values, dtype_hint=tf.float32, name='constant_values')
    142 min_event_ndims_ = int(-np.min(np.pad(
--> [143](https://file+.vscode-resource.vscode-cdn.net/home/robbin/Projects/pytsa/notebooks/~/Projects/pytsa/.venv/lib/python3.12/site-packages/tensorflow_probability/substrates/jax/bijectors/pad.py:143)     np.reshape(axis_, newshape=[-1]),
    144     mode='constant', pad_width=[[0, 1]])))
    145 super(Pad, self).__init__(
    146     forward_min_event_ndims=min_event_ndims_,
    147     inverse_min_event_ndims=min_event_ndims_,
   (...)    150     parameters=parameters,
    151     name=name)

TypeError: reshape() got an unexpected keyword argument 'newshape'

Update:
It seems to derive from: .venv/lib/python3.12/site-packages/tensorflow_probability/substrates/jax/bijectors/pad.py:143
My guess is that the function does not have a newshape argument but rather a shape argument.

[rmndrs89]

It seems that the issue should have been solved in a newer commit in the tensorflow_probability library, see: here

but I am not sure if this commit has already been included in the latest release.

[retconrecon]

This is the same root cause as #428 — stable tensorflow-probability (0.25.0) uses deprecated APIs that break with recent JAX/NumPy versions. The newshape kwarg in pad.py:143 was removed in NumPy 2.0.

The current tfp-nightly dependency in main fixes this, and I've opened a PR to pin it to a verified build (==0.26.0.dev20260101) — tested against JAX 0.8.1 + NumPy 2.3.5 with GaussianHMM initialization passing.