class_docstr: Since tfd.Sample uses transposes, it requires a custom event bijector.
class_docstr_tok: ["Since","tfd.Sample","uses","transposes",",","it","requires","a","custom","event","bijector","."]
code: def _transpose_around_bijector_fn(
      self,
      bijector_fn,
      arg,
      src_event_ndims,
      dest_event_ndims=None,
      fn_reduces_event=False,
      **kwargs,
  ):
      # This function moves the axes corresponding to `self.sample_shape` to the
      # left of the batch shape, then applies `bijector_fn`, then moves the axes
      # corresponding to `self.sample_shape` back to the event part of the shape.
      #
      # `src_event_ndims` and `dest_event_ndims` indicate the expected event rank
      # (omitting `self.sample_shape`) before and after applying `bijector_fn`.
      #
      # This function arose because forward and inverse ended up being quite
      # similar. It was then only a small generalization to also support {F/I}LDJ.
      batch_ndims = ps.rank_from_shape(
          self.distribution.batch_shape_tensor,
          self.distribution.batch_shape,
      )
      extra_sample_ndims = ps.rank_from_shape(self.sample_shape)
      arg_ndims = ps.rank(arg)
      # (1) Expand arg's dims.
      d = arg_ndims - batch_ndims - extra_sample_ndims - src_event_ndims
      arg = tf.reshape(
          arg,
          shape=ps.pad(
              ps.shape(arg),
              paddings=[[ps.maximum(0, -d), 0]],
              constant_values=1,
          ),
      )
      arg_ndims = ps.rank(arg)
      sample_ndims = ps.maximum(0, d)
      # (2) Transpose arg's dims.
      sample_dims = ps.range(0, sample_ndims)
      batch_dims = ps.range(sample_ndims, sample_ndims + batch_ndims)
      extra_sample_dims = ps.range(
          sample_ndims + batch_ndims,
          sample_ndims + batch_ndims + extra_sample_ndims,
      )
      event_dims = ps.range(
          sample_ndims + batch_ndims + extra_sample_ndims, arg_ndims
      )
      perm = ps.concat(
          [sample_dims, extra_sample_dims, batch_dims, event_dims],
          axis=0,
      )
      arg = tf.transpose(arg, perm=perm)
      # (3) Apply underlying bijector.
      result = bijector_fn(arg, **kwargs)
      # (4) Transpose sample_shape from the sample to the event shape.
      result_ndims = ps.rank(result)
      if fn_reduces_event:
          dest_event_ndims = 0
      d = (
          result_ndims
          - batch_ndims
          - extra_sample_ndims
          - dest_event_ndims
      )
      if fn_reduces_event:
          # In some cases, fn may reduce event too far, i.e. ildj may return a
          # scalar `0.`, which won't work with the transpose we do below.
          result = tf.reshape(
              result,
              shape=ps.pad(
                  ps.shape(result),
                  paddings=[[ps.maximum(0, -d), 0]],
                  constant_values=1,
              ),
          )
          result_ndims = ps.rank(result)
      sample_ndims = ps.maximum(0, d)
      sample_dims = ps.range(0, sample_ndims)
      extra_sample_dims = ps.range(
          sample_ndims, sample_ndims + extra_sample_ndims
      )
      batch_dims = ps.range(
          sample_ndims + extra_sample_ndims,
          sample_ndims + extra_sample_ndims + batch_ndims,
      )
      event_dims = ps.range(
          sample_ndims + extra_sample_ndims + batch_ndims, result_ndims
      )
      perm = ps.concat(
          [sample_dims, batch_dims, extra_sample_dims, event_dims],
          axis=0,
      )
      return tf.transpose(result, perm=perm)
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  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 dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 6 for calling the _NonCompositeTensorDefaultSampleBijector obj's underlying member method code with example usage: obj._transpose_around_bijector_fn(bijector_fn, arg, src_event_ndims, dest_event_ndims, fn_reduces_event) and returns: tf