class_docstr: A JDMixin that shards the log_prob calculation.
docstr: Executes a model, adding `pbroadcasts` to ensure correct gradients.
docstr_tok: ["Executes","a","model",",","adding","`","pbroadcasts","`","to","ensure","correct","gradients","."]
code: def _distribute_execute_model(
      self,
      sample_shape=(),
      seed=None,
      value=None,
      sample_and_trace_fn=jd_lib.trace_distributions_and_values,
  ):
      """Executes a model, adding `pbroadcasts` to ensure correct gradients."""
      shard_axis_names = self._model_flatten(
          self.experimental_shard_axis_names
      )
      final_values_out = []
      if value is None:
          value = ()
  
      def sample_and_trace_value_fn(
          dist, sample_shape, seed, value=None
      ):
          value, traced = sample_and_trace_fn(
              dist=dist,
              sample_shape=sample_shape,
              seed=seed,
              value=value,
          )
          # We trace `next_value` here so we can pass it back in as part of `value`
          # in the next iteration of the coroutine.
          return value, (value, traced)
  
      for output_index, output_axes in enumerate(shard_axis_names):
          # We pbroadcast all values according to the difference between the current
          # `output_axes` and their own active axes.
          previous_shard_axes = shard_axis_names[: len(value)]
          pbroadcasted_value = tuple(
              pbroadcast_value(v, v_axis_names, output_axes)
              for v, v_axis_names in zip(value, previous_shard_axes)
          )
          pbroadcasted_values, traced_values = zip(
              *super()._execute_model(
                  sample_shape=sample_shape,
                  seed=seed,
                  value=pbroadcasted_value + (None,),
                  stop_index=output_index + 1,
                  sample_and_trace_fn=sample_and_trace_value_fn,
              )
          )
          value = _maybe_substitute_or_add_value_in_tuple(
              value, output_index, pbroadcasted_values[output_index]
          )
          final_values_out.append(traced_values[output_index])
      return final_values_out
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.distributions import joint_distribution
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import samplers
label_desc: Use image node_id 2 for calling the JointDistributionDistributedMixin obj's underlying member method code with example usage: obj._distribute_execute_model(sample_shape, seed, value, sample_and_trace_fn) and returns: final_values_out, value,  value,  traced