code: def forward(
      self,
      text: torch.Tensor,
      text_len: torch.Tensor,
      return_probe_features: bool = False,
  ) -> torch.Tensor:
      if len(text.shape) != 3:
          raise ValueError("expected text to be 3d")
      if len(text_len.shape) != 1:
          raise ValueError("expected text length to be 1d")
      if text.shape[0] != text_len.shape[0]:
          raise ValueError(
              "text and text_len have inconsistent batch dimensions"
          )
  
      index = (text_len - 1)[:, None, None].expand(-1, 1, text.shape[2])
      x = torch.gather(text, dim=1, index=index)
      assert list(x.shape) == [text.shape[0], 1, text.shape[2]]
  
      if return_probe_features:
          return x[:, 0]
  
      x = self.ln(x)
      return self.f(x[:, 0])
  
imports: import math
  from collections import OrderedDict
  from typing import List, Optional, Tuple, cast
  import attr
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from .attention import AttentionInfo, DenseAttentionMask, DenseCausalAttentionMask, make_full_layout, to_attention_info
  from .utils import Affine, LayerNorm, zero_key_bias_grad
label_desc: Use image node_id 2 for calling the TextFeatureExtractor obj's underlying member method code with example usage: obj.forward(text, text_len, return_probe_features) and returns: self, x