Relational DB in Jax on GPU/TPU only?

[MRiabov]

Sup Jax,
I'm writing a model in Jax - it should be very performant on GPUs and TPUs, and as everyone knows using GPU-only scripts with vectorization can speed up the app by three orders of magnitude.

So, I have a relational database with 1000 entries that my model would use during runtime - are there any databases that support jax's vmap? For everything else is already in tensors.

I hope that somebody from the community would know.

Thanks everyone.

P.S. I'm coming from traditional SWE, and we use Relational DBs for almost everything. Maybe I could use something else, but I need an O(1) fetch complexity (storing is, I believe, rather unlikely - it will probably be in other arrays.

P.S.S. Was also posted on https://www.reddit.com/r/MLQuestions/comments/1dsrnrl/relational_db_on_gputpu_only/