tinybox?

Not sure how suitable this is for you.

https://tinygrad.org/#tinybox