class_docstr: A classifier that uses attention to summarize the inputs.
  
  Attributes:
    hidden_size: The hidden size of the MLP classifier.
    output_size: The number of output classes for the classifier.
    dropout_rate: The dropout rate applied over the encoded_inputs, the summary
      of the inputs, and inside the MLP. Applied when `deterministic` is False.
    deterministic: Disables dropout if True.
docstr: Applies model to the encoded inputs.
  
  Args:
    encoded_inputs: The inputs (e.g., sentences) that have already been
      encoded by some encoder, e.g., an LSTM. <float32>[batch_size,
      seq_length, encoded_inputs_size].
    lengths: The lengths of the inputs. <int64>[batch_size].
    deterministic: Disables dropout when set to True.
  
  Returns:
    An array of logits <float32>[batch_size, output_size].
docstr_tok: ["Applies","model","to","the","encoded","inputs",".","Args",":","encoded_inputs",":","The","inputs","(","e.g.",",","sentences",")","that","have","already","been","encoded","by","some","encoder",",","e.g.",",","an","LSTM",".","<","float32",">","[","batch_size",",","seq_length",",","encoded_inputs_size","]",".","lengths",":","The","lengths","of","the","inputs",".","<","int64",">","[","batch_size","]",".","deterministic",":","Disables","dropout","when","set","to","True",".","Returns",":","An","array","of","logits","<","float32",">","[","batch_size",",","output_size","]","."]
code: def __call__(
      self,
      encoded_inputs: Array,
      lengths: Array,
      deterministic: bool | None = None,
  ) -> Array:
      """Applies model to the encoded inputs.
  
      Args:
        encoded_inputs: The inputs (e.g., sentences) that have already been
          encoded by some encoder, e.g., an LSTM. <float32>[batch_size,
          seq_length, encoded_inputs_size].
        lengths: The lengths of the inputs. <int64>[batch_size].
        deterministic: Disables dropout when set to True.
  
      Returns:
        An array of logits <float32>[batch_size, output_size].
      """
      deterministic = nn.module.merge_param(
          "deterministic", self.deterministic, deterministic
      )
      encoded_inputs = self.dropout_layer(
          encoded_inputs, deterministic=deterministic
      )
  
      # Compute attention. attention.shape: <float32>[batch_size, seq_len].
      mask = sequence_mask(lengths, encoded_inputs.shape[1])
      attention = self.keys_only_mlp_attention(encoded_inputs, mask)
  
      # Summarize the inputs by taking their weighted sum using attention scores.
      context = jnp.expand_dims(attention, 1) @ encoded_inputs
      context = context.squeeze(
          1
      )  # <float32>[batch_size, encoded_inputs_size]
      context = self.dropout_layer(context, deterministic=deterministic)
  
      # Make the final prediction from the context vector (the summarized inputs).
      logits = self.mlp(context, deterministic=deterministic)
      return logits
  
imports: from __future__ import annotations
  import functools
  from typing import Any, Callable, Optional
  from flax import linen
  import jax
  from jax import numpy
  from jax._src.typing import Array
label_desc: Use image node_id 2 for calling the AttentionClassifier obj's underlying member method code with example usage: obj.__call__(encoded_inputs, lengths, deterministic) and returns: logits