Could RIAB support GPU via jax (import jax.numpy as jnp)
This should be backward compatible. I.e. a users should optionally specify a GPU usage flag otherwise numpy is used as normal.
GPU would not massively speed up RIAB except for FeedForwardCells and any use-case where synaptic weight matrices into FeedForwardCells are learnt, typically requiring large N_cell x N_cells matrix multiplications. I.e it is likely that non-GPU will continue to satisfy a large majority of use cases.
For now I do not intend to use jax to pre-compile or vectorise RIAB code, a change which would require significant and likely backwards incompatible modifications to the code base.
Could RIAB support GPU via jax (
import jax.numpy as jnp
)This should be backward compatible. I.e. a users should optionally specify a GPU usage flag otherwise
numpy
is used as normal.GPU would not massively speed up RIAB except for
FeedForwardCells
and any use-case where synaptic weight matrices into FeedForwardCells are learnt, typically requiring large N_cell x N_cells matrix multiplications. I.e it is likely that non-GPU will continue to satisfy a large majority of use cases.For now I do not intend to use jax to pre-compile or vectorise RIAB code, a change which would require significant and likely backwards incompatible modifications to the code base.