docstr: Generate and return predictions, and randomly choose a prediction weighted by the scores
  
  TODO: pass in a temperature
code: def weighted_choice(self, states):
      """
      Generate and return predictions, and randomly choose a prediction weighted by the scores
  
      TODO: pass in a temperature
      """
      predictions = self.forward(states)
      pred_trans = []
      all_scores = []
      for state, prediction in zip(states, predictions):
          legal_idx = [
              idx
              for idx in range(prediction.shape[0])
              if self.transitions[idx].is_legal(state, self)
          ]
          if len(legal_idx) == 0:
              pred_trans.append(None)
              continue
          scores = prediction[legal_idx]
          scores = torch.softmax(scores, dim=0)
          idx = torch.multinomial(scores, 1)
          idx = legal_idx[idx]
          pred_trans.append(self.transitions[idx])
          all_scores.append(prediction[idx])
      all_scores = torch.stack(all_scores)
      return predictions, pred_trans, all_scores
  
imports: from collections import namedtuple
  from enum import Enum
  import logging
  import math
  import random
  import torch
  import torch.nn
  from torch.nn.utils.rnn import pack_padded_sequence
  from stanza.models.common.bert_embedding import extract_bert_embeddings
  from stanza.models.common.maxout_linear import MaxoutLinear
  from stanza.models.common.utils import unsort
  from stanza.models.common.vocab import PAD_ID, UNK_ID
  from stanza.models.constituency.base_model import BaseModel
  from stanza.models.constituency.label_attention import LabelAttentionModule
  from stanza.models.constituency.lstm_tree_stack import LSTMTreeStack
  from stanza.models.constituency.parse_transitions import TransitionScheme
  from stanza.models.constituency.parse_tree import Tree
  from stanza.models.constituency.partitioned_transformer import PartitionedTransformerModule
  from stanza.models.constituency.positional_encoding import ConcatSinusoidalEncoding
  from stanza.models.constituency.transformer_tree_stack import TransformerTreeStack
  from stanza.models.constituency.tree_stack import TreeStack
  from stanza.models.constituency.utils import build_nonlinearity, initialize_linear
label_desc: Use image node_id 28 for calling the LSTMModel obj's underlying member method code with example usage: obj.weighted_choice(states) and returns: predictions,  pred_trans,  all_scores