Simplify `B.mean` for PyTorch (closes #21)

Fix support for custom JAX operations

Don't force tests to run with NumPy 2

Fix `device` for JAX

Add fdm as test dependency

Fix `keras.KerasTensor` availability issue

Fix Python 3.8

Merge branch 'master' of github.com:wesselb/lab

Update README.md

Merge pull request #14 from rossviljoen/rv/cast-type-fix

small fix to cast type signature