tfp.experimental.math.patch_manual_special_functions

Patches in the manually implemented special functions.

@contextlib.contextmanager
tfp.experimental.math.patch_manual_special_functions()

Normally you'd just use functions coming from the array library you're using, but on some platforms (like TPU) the default implementations are insufficiently precise for certain tasks when running under 32 bits (64 bit implementations are typically okay).

This patches in manual implementations of those functions which are provide higher precision at the cost of speed. The list of affected functions is:

  • exp
  • log
  • expm1
  • log1p
  • logsumexp (aka reduce_logsumexp)
  • softplus

Nothing.