jaxtyping VS tiny-cuda-nn

You want runtime typechecking.

See either beartype [1] or typeguard [2]. And if you're doing any kind of array-based programming (JAX or not), then jaxtyping [3].

[1] https://github.com/beartype/beartype/

[2] https://github.com/agronholm/typeguard

[3] https://github.com/patrick-kidger/jaxtyping

Try using [jaxtyping](https://github.com/google/jaxtyping).

It also supports numpy/pytorch/etc.

Since you mention ML use-cases, you might like jaxtyping.

jaxtyping: rich shape & dtype annotations for arrays and tensors (also supports PyTorch/TensorFlow/NumPy);

Heads-up that my newer jaxtyping project now exists.

As other folks have commented, type hints are now a big deal. For static typing the best checker is pyright. For runtime checking there is typeguard and beartype. These can be integrated with array libraries through jaxtyping. (Which also works for PyTorch/numpy/etc., despite the name.)

In the opposite direction from your question is a very interesting project, TinyNN all implemented as close to the metal as possible and very fast: https://github.com/NVlabs/tiny-cuda-nn

This repo only implemented the rendering part of the NGP but is more simple and has a lesser amount of code compared to the original (Instant-NGP and tiny-cuda-nn).

...a C++ library with a CUDA backend. But these high-performance building blocks might only be saturating the GPU fully if the data is large enough.

I haven't looked at implementing these things, but I imagine uf you have smaller networks and thus less data, the large building blocks may not be optimal. You may for example want to fuse some operations to reduce memory latency from repeated memory access.

In PyTorch world, there are approaches for small networks as well, there is https://github.com/NVlabs/tiny-cuda-nn - as far as I understand from the first link in the README, it makes clever use of the CUDA shared memory, which can hold all the weights of a tiny network (but not larger ones).