code: def batched_span_select(
      target: torch.Tensor, spans: torch.LongTensor
  ) -> torch.Tensor:
      """
      The given `spans` of size `(batch_size, num_spans, 2)` indexes into the sequence
      dimension (dimension 2) of the target, which has size `(batch_size, sequence_length,
      embedding_size)`.
      This function returns segmented spans in the target with respect to the provided span indices.
      It does not guarantee element order within each span.
      # Parameters
      target : `torch.Tensor`, required.
          A 3 dimensional tensor of shape (batch_size, sequence_length, embedding_size).
          This is the tensor to be indexed.
      indices : `torch.LongTensor`
          A 3 dimensional tensor of shape (batch_size, num_spans, 2) representing start and end
          indices (both inclusive) into the `sequence_length` dimension of the `target` tensor.
      # Returns
      span_embeddings : `torch.Tensor`
          A tensor with shape (batch_size, num_spans, max_batch_span_width, embedding_size]
          representing the embedded spans extracted from the batch flattened target tensor.
      span_mask: `torch.BoolTensor`
          A tensor with shape (batch_size, num_spans, max_batch_span_width) representing the mask on
          the returned span embeddings.
      """
      # both of shape (batch_size, num_spans, 1)
      span_starts, span_ends = spans.split(1, dim=-1)
  
      # shape (batch_size, num_spans, 1)
      # These span widths are off by 1, because the span ends are `inclusive`.
      span_widths = span_ends - span_starts
  
      # We need to know the maximum span width so we can
      # generate indices to extract the spans from the sequence tensor.
      # These indices will then get masked below, such that if the length
      # of a given span is smaller than the max, the rest of the values
      # are masked.
      max_batch_span_width = span_widths.max().item() + 1
  
      # Shape: (1, 1, max_batch_span_width)
      max_span_range_indices = get_range_vector(
          max_batch_span_width, get_device_of(target)
      ).view(1, 1, -1)
      # Shape: (batch_size, num_spans, max_batch_span_width)
      # This is a broadcasted comparison - for each span we are considering,
      # we are creating a range vector of size max_span_width, but masking values
      # which are greater than the actual length of the span.
      #
      # We're using <= here (and for the mask below) because the span ends are
      # inclusive, so we want to include indices which are equal to span_widths rather
      # than using it as a non-inclusive upper bound.
      span_mask = max_span_range_indices <= span_widths
      raw_span_indices = span_ends - max_span_range_indices
      # We also don't want to include span indices which are less than zero,
      # which happens because some spans near the beginning of the sequence
      # have an end index < max_batch_span_width, so we add this to the mask here.
      span_mask = span_mask & (raw_span_indices >= 0)
      span_indices = torch.nn.functional.relu(
          raw_span_indices.float()
      ).long()
  
      # Shape: (batch_size, num_spans, max_batch_span_width, embedding_dim)
      span_embeddings = batched_index_select(target, span_indices)
  
      return span_embeddings, span_mask
  
imports: from typing import Any, Dict, List, Optional, Sequence, Tuple, TypeVar, Union
  import torch
label_desc: Use image node_id 5 for calling a global function with example usage: batched_span_select(target, spans) and returns: span_embeddings,  span_mask