code: def einsum_backward_sample(operands, sample1, sample2):
      """
      Cuts down samples to pass on to subsequent steps.
      This is used in various ``_EinsumBackward.__call__()`` methods.
      This assumes all operands have a ``._pyro_dims`` attribute set.
      """
      # Combine upstream sample with sample at this site.
      if sample1 is None:
          sample = sample2
      elif sample2 is None:
          sample = sample1
      else:
          # Slice sample1 down based on choices in sample2.
          assert set(sample1._pyro_sample_dims).isdisjoint(
              sample2._pyro_sample_dims
          )
          sample_dims = (
              sample1._pyro_sample_dims + sample2._pyro_sample_dims
          )
          for dim, index in zip(
              sample2._pyro_sample_dims, jit_iter(sample2)
          ):
              if dim in sample1._pyro_dims:
                  index._pyro_dims = sample2._pyro_dims[1:]
                  sample1 = packed.gather(sample1, index, dim)
  
          # Concatenate the two samples.
          parts = packed.broadcast_all(sample1, sample2)
          sample = torch.cat(parts)
          sample._pyro_dims = parts[0]._pyro_dims
          sample._pyro_sample_dims = sample_dims
          assert sample.dim() == len(sample._pyro_dims)
          if not torch._C._get_tracing_state():
              assert sample.size(0) == len(sample._pyro_sample_dims)
  
      # Select sample dimensions to pass on to downstream sites.
      for x in operands:
          if not hasattr(x, "_pyro_backward"):
              continue
          if sample is None:
              yield x._pyro_backward, None
              continue
          x_sample_dims = set(x._pyro_dims) & set(
              sample._pyro_sample_dims
          )
          if not x_sample_dims:
              yield x._pyro_backward, None
              continue
          if x_sample_dims == set(sample._pyro_sample_dims):
              yield x._pyro_backward, sample
              continue
          x_sample_dims = "".join(sorted(x_sample_dims))
          x_sample = sample[
              [
                  sample._pyro_sample_dims.index(dim)
                  for dim in x_sample_dims
              ]
          ]
          x_sample._pyro_dims = sample._pyro_dims
          x_sample._pyro_sample_dims = x_sample_dims
          assert x_sample.dim() == len(x_sample._pyro_dims)
          if not torch._C._get_tracing_state():
              assert x_sample.size(0) == len(x_sample._pyro_sample_dims)
          yield x._pyro_backward, x_sample
  
imports: import weakref
  from abc import ABCMeta, abstractmethod
  import torch
  from pyro.ops import packed
  from pyro.util import jit_iter
label_desc: Use image node_id 3 for calling a global function with example usage: einsum_backward_sample(operands, sample1, sample2) without return types