docstr: args is a list of tensors, to be passed to self.likelihoods.<func_name>
  
  args[-1] is the 'Y' argument, which contains the indexes to self.likelihoods.
  
  This function splits up the args using dynamic_partition, calls the
  relevant function on the likelihoods, and re-combines the result.
docstr_tok: ["args","is","a","list","of","tensors",",","to","be","passed","to","self.likelihoods.","<","func_name",">","args","[","-1","]","is","the","'","Y","'","argument",",","which","contains","the","indexes","to","self.likelihoods",".","This","function","splits","up","the","args","using","dynamic_partition",",","calls","the","relevant","function","on","the","likelihoods",",","and","re-combines","the","result","."]
code: def _partition_and_stitch(
      self, args: Sequence[TensorType], func_name: str
  ) -> tf.Tensor:
      """
      args is a list of tensors, to be passed to self.likelihoods.<func_name>
  
      args[-1] is the 'Y' argument, which contains the indexes to self.likelihoods.
  
      This function splits up the args using dynamic_partition, calls the
      relevant function on the likelihoods, and re-combines the result.
      """
      # get the index from Y
      args_list = list(args)
      Y = args_list[-1]
      ind = Y[..., -1]
      ind = tf.cast(ind, tf.int32)
      Y = Y[..., :-1]
      args_list[-1] = Y
  
      # split up the arguments into chunks corresponding to the relevant likelihoods
      args_chunks = zip(
          *[
              tf.dynamic_partition(X, ind, len(self.likelihoods))
              for X in args_list
          ]
      )
  
      # apply the likelihood-function to each section of the data
      funcs = [getattr(lik, func_name) for lik in self.likelihoods]
      results = [f(*args_i) for f, args_i in zip(funcs, args_chunks)]
  
      # stitch the results back together
      partitions = tf.dynamic_partition(
          tf.range(0, tf.size(ind)), ind, len(self.likelihoods)
      )
      results = tf.dynamic_stitch(partitions, results)
  
      return results
  
imports: import abc
  from typing import Any, Callable, Iterable, Optional, Sequence, Union
  import tensorflow
  from check_shapes import check_shapes, inherit_check_shapes
  from ..base import MeanAndVariance, Module, TensorType
  from ..quadrature import GaussianQuadrature, NDiagGHQuadrature, ndiag_mc
label_desc: Use image node_id 2 for calling the SwitchedLikelihood obj's underlying member method code with example usage: obj._partition_and_stitch(args, func_name) and returns: results