Make tfd.Gamma.sample use log_space sampling under XLA/JAX.

Benchmarks have shown that log_space sampling is a bit slower in Graph mode, so
we keep the old behavior for that configuration.

This should help JAX the most, which typically does not have 64 bit dtype
enabled. The old default assumed it was, causing warnings and reduced numerical
precision.

PiperOrigin-RevId: 452391450

Author: SiegeLordEx

tensorflow_probability/python/distributions/gamma.py

@@ -265,14 +265,19 @@ def _log_rate_parameter(self):
       caveats.""")
   def _sample_n(self, n, seed=None):
     seed = samplers.sanitize_seed(seed, salt='gamma')
+    log_space = implementation_selection.is_xla()
 
-    return random_gamma(
+    res = random_gamma(
         shape=ps.convert_to_shape_tensor([n]),
         concentration=tf.convert_to_tensor(self.concentration),
         rate=None if self.rate is None else tf.convert_to_tensor(self.rate),
         log_rate=(None if self.log_rate is None else
                   tf.convert_to_tensor(self.log_rate)),
+        log_space=log_space,
         seed=seed)
+    if log_space:
+      res = tf.math.exp(res)
+    return res
 
   def _log_prob(self, x, rate=None):
     concentration = tf.convert_to_tensor(self.concentration)

tensorflow_probability/python/internal/implementation_selection.py

@@ -22,6 +22,7 @@
 
 __all__ = [
     'implementation_selecting',
+    'is_xla',
     'never_runs_functions_eagerly',
 ]
 
@@ -47,7 +48,7 @@
 NUMPY_MODE = False
 
 
-def _is_xla():
+def is_xla():
   """Returns `True` when we are tracing a function for XLA compilation."""
   if JAX_MODE:
     return True
@@ -134,7 +135,7 @@ def stub_fn(**kwargs):
 
   def impl_selecting_fn(**kwargs):
     """The wrapper function to be returned."""
-    if _is_xla():  # JAX, XLA breakout.
+    if is_xla():  # JAX, XLA breakout.
       return default_fn(**kwargs)
     if NUMPY_MODE:  # Numpy breakout.
       return cpu_fn(**kwargs)