Mixed precision / autocast for bf16 and fp16 on NKI

[scttfrdmn]

Support dtype="bf16" / dtype="fp16" overrides on einsum. Align with Neuron SDK's autocast recommendations for Trainium.

[scttfrdmn]

Implemented in 25de104: dtype= kwarg on einsum() with string aliases and torch.dtype support.