code: def _make_masked_fn(
      fn_name, n_event_shapes, safe_value, make_arg0_safe=False
  ):
      """Implements functions like mean, variance, etc.
  
      Args:
        fn_name: Name of the method called on the underlying distribution.
        n_event_shapes: Number of event shape repeats in the shape of the underlying
          function's output.
        safe_value: The value to be placed in invalid locations. May be
          `'safe_sample'` to specify we should use the "safe sample" value.
        make_arg0_safe: If `True`, we will apply `self.safe_sample_fn` to ensure the
          argument passed into the underlying routine is a "safe" sample.
  
      Returns:
        fn: Callable implementing the given function.
      """
  
      def fn(self, *args, **kwargs):
          if (
              safe_value == "safe_sample" or make_arg0_safe
          ):  # Only if needed.
              safe_val = tf.stop_gradient(
                  self.safe_sample_fn(self.distribution)
              )
  
          validity_mask = tf.convert_to_tensor(self.validity_mask)
          if make_arg0_safe:
              x = args[0]
              safe_x = tf.where(
                  _add_event_dims_to_mask(validity_mask, dist=self),
                  x,
                  safe_val,
              )
              args = (safe_x,) + args[1:]
  
          val = getattr(self.distribution, fn_name)(*args, **kwargs)
          if n_event_shapes:
              validity_mask = tf.reshape(
                  validity_mask,
                  ps.concat(
                      [ps.shape(validity_mask)]
                      + [ps.ones_like(self.event_shape_tensor())]
                      * n_event_shapes,
                      axis=0,
                  ),
              )
          if safe_value == "safe_sample":
              sentinel = tf.cast(safe_val, val.dtype)
          else:
              sentinel = tf.cast(safe_value, val.dtype)
          return tf.where(validity_mask, val, sentinel)
  
      fn.__name__ = f"_{fn_name}"
      return fn
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.distributions import batch_broadcast
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
label_desc: Use image node_id 2 for calling a global function with example usage: _make_masked_fn(fn_name, n_event_shapes, safe_value, make_arg0_safe) and returns: fn, tf