code: def _chunked_selector_output_shape(  # pylint: disable=invalid-name
      input_shapes, selector=None, **unused_kwargs):
    
    # Read the main function below first, the shape logic just follows the ops.
    selector = selector or (lambda x: [] if x < 1 else [x-1])
    triples, _ = zip(*input_shapes)
    (query_shapes, key_shapes, value_shapes) = zip(*triples)
    result = []
    for i in range(len(input_shapes)):
      selected = selector(i)
      cur_key_shape, cur_value_shape = key_shapes[i], value_shapes[i]
      # Since keys and values are [batch, length, depth] we concatenate on axis=1.
      new_key_len = sum([key_shapes[j][1] for j in selected]) + cur_key_shape[1]
      new_key_shape = (cur_key_shape[0], new_key_len, cur_key_shape[2])
      new_value_len = sum(
          [value_shapes[j][1] for j in selected]) + cur_value_shape[1]
      new_value_shape = (cur_value_shape[0], new_value_len, cur_value_shape[2])
      # Masks are (1, query-len, key-len).
      new_mask_shape = (1, query_shapes[i][1], new_key_len)
      new_shape = ((query_shapes[i], new_key_shape, new_value_shape),
                   new_mask_shape)
      result.append(new_shape)
    return tuple(result)
docs: Helper: calculate output shape for chunked key selector (see below).