Mixed precision training support

[ogencoglu]

Does compute_gramian() support mixed precision training? I couldn't make it work at least.

[PierreQuinton]

Thanks for the issue!

It is not clear to me what the problem/goal is. Could you provide us with a minimal example of what you want to do?

[ValerianRey]

Had a bug a few days ago with autojac.jac related to bfloat16. It wasn't autogram but it still may be related. I'll try to post a simple reproducer.

[ValerianRey]

The bug I had with AMP is documented here: #699
It's very possible that this bug also affects autogram.Engine.compute_gramian, and not just autojac (we use vmap quite similarly internally). So it's possible that the two issues are related. I'll try to fix that soon.

@ogencoglu do you remember what went wrong when you tried compute_gramian on automatic mixed precision? Could it be the same error?

[ogencoglu]

Unfortunately, I don't remember.