Add mixed precision

[seanmor5]

WIP.

Need to think of a better way to test

[seanmor5]

For reference, this is based off of: https://github.com/deepmind/jmp

[josevalim]

LGTM! Although I would consider working on some lightweight tests although I am not quite sure how to test this... maybe by testing directly on the Nx expression?

[seanmor5]

@josevalim I think that might be the best way also, the expression might end up being really large. There is also a problem right now where I don't think it is respecting the policy in all cases. I can guarantee the policy is correct for compute and output because it's just a simple cast; however, the params policy doesn't seem to stay consistent even though they're initialized to {:f, 16} or {:bf, 16}