+1 for JAX. Basically designed to be the successor to TensorFlow, and much nicer to work with. Strangely I've not seen it discussed around HN much but it's what I do 100% of my work in these days.

Whilst I'm here: shameless self-promotion for Equinox and Diffrax:

https://github.com/patrick-kidger/equinox https://github.com/patrick-kidger/diffrax

Which are neural network and differential equation libraries for JAX.

[Obligatory I-am-a-googler-my-opinions-do-not-represent-your-employer...]