Add context manager api for setting mixed precision policies.

[copybara-service[bot]]

Add context manager api for setting mixed precision policies.

mod = hk.Linear(1) with hk.mixed_precision.push_policy(hk.Linear, policy): ... x = mod(x) # policy is active ... x = mod(x) # previous policy (if any) is active