API design question - should users ever be directly exposed to JAX?

[Arun-Niranjan]

There is a case to be made for abstracting out JAX from the pymdp API (I got this feedback from a user at IWAI).

This would have the following benefits:

1. plenty of python users are familiar with numpy, but not JAX
2. in a tutorial setting, it would be much easier to demonstrate the improvements JAX brings with minimal migration effort of existing code
3. if we ever decide to move away from JAX and use something else, we could do it under the hood without any breaking changes
4. it would increase the flexibility we can offer in terms of "bring your own" sampler/inference algorithm as per More flexible configuration of inference algorithms #275

However I do have concerns:

1. JAX already has jax.numpy which is close enough to existing numpy for most python users' needs
2. JAX works on CPU as well as GPU and TPU, so once you have your JAX arrays for A, B, C etc. I don't believe anything would need to change anyway
3. JAX at this point is pretty intermeshed in our codebase, it would probably be less effort to just create a new package using a different technology if we wanted to. (also how would we handle return types? converting them all to custom types would likely have a performance impact)
4. This is the most exciting reason for this idea to me - but in practice I think it would be difficult to establish the correct abstraction without knowing more about how users would want this feature. also validation that a bring your own function is JAX compatible could be trickier if it isn't using JAX types and primitives

I am leaning 70/30 away from this idea, but I wanted to document it here and get opinions. If someone can address the concerns and thinks the reward is worthwhile to do the work before we do a v1 release then we can prioritise it accordingly.

[conorheins]

I do agree with you, and think we should not hide users from JAX at this stage, this seems like more of a pymdp 2.0.0 thing. As a result I've moved this off the to-do list for the 1.0.0 release roadmap