Use np.array for typing, remove parameter unfreeze

lindermanlab / S5

MIT License

248 stars 43 forks source link

Use np.array for typing, remove parameter unfreeze #9

Closed davidzoltowski closed 11 months ago

davidzoltowski commented 1 year ago

This PR makes two small changes to work with newer versions of Jax and Flax.

The typing in ssm.py is changed to np.array as np.DeviceArray is deprecated
Remove the unfreeze() call in train_helper.py

davidzoltowski commented 11 months ago

Closing this PR to stop new development on our branch from getting merged into this PR. We can open a separate branch and PR to resubmit the changes from the first commit of this PR.