class_docstr: Seq2seq task: encoder's input is "x+y", decoder's output is "=z".
class_docstr_tok: ["Seq2seq","task",":","encoder","'s","input","is","``","x+y","''",",","decoder","'s","output","is","``","=z","''","."]
code: def __init__(
      self,
      batch_size: int = 128,
      max_len_query_digit: int = 3,
      test: bool = False,
  ):
      char_table = CharacterTable()
      max_input_len = max_len_query_digit + 2 + 2
      max_output_len = max_len_query_digit + 3
      max_num = pow(10, max_len_query_digit)
      self.obs_shape = tuple([max_input_len, char_table.vocab_size])
      self.act_shape = tuple([max_output_len, char_table.vocab_size])
      self.max_steps = 1
  
      def encode_onehot(batch_inputs, max_len):
          def encode_str(s):
              tokens = char_table.encode(s)
              org_len = len(tokens)
              tokens = jnp.pad(
                  tokens, [(0, max_len - org_len)], mode="constant"
              )
              return jax.nn.one_hot(
                  tokens, char_table.vocab_size, dtype=jnp.float32
              )
  
          return jnp.array([encode_str(inp) for inp in batch_inputs])
  
      def decode_onehot(batch_inputs):
          return np.array(
              list(
                  map(
                      lambda x: char_table.decode(x.argmax(axis=-1)),
                      batch_inputs,
                  )
              )
          )
  
      self.decode_embeddings = decode_onehot
  
      def breakdown_int(n):
          return jnp.where(
              n == 0,
              jnp.zeros(max_len_query_digit),
              jnp.array(
                  [
                      (n // (10**i)) % 10
                      for i in range(max_len_query_digit - 1, -1, -1)
                  ]
              ),
          )
  
      def get_batch_data(key):
          keys = random.split(key, 2)
          add_op1 = random.randint(
              keys[0], minval=0, maxval=100, shape=(batch_size,)
          )
          add_op2 = random.randint(
              keys[1], minval=0, maxval=max_num, shape=(batch_size,)
          )
          for op1, op2 in zip(add_op1, add_op2):
              inputs = jnp.concatenate(
                  [
                      breakdown_int(op1)[1:],
                      jnp.array(
                          [
                              10,
                          ]
                      ),  # 10 is '+'
                      breakdown_int(op2),
                  ],
                  axis=0,
              )
              outputs = jnp.concatenate(
                  [
                      jnp.array(
                          [
                              11,
                          ]
                      ),  # 11 is '='
                      breakdown_int(op1 + op2),
                  ],
                  axis=0,
              )
              yield inputs, outputs
  
      def reset_fn(key):
          inputs, outputs = zip(*get_batch_data(key[0]))
          batch_data = encode_onehot(inputs, max_input_len)
          batch_labels = encode_onehot(outputs, max_output_len)
          return State(
              obs=jnp.repeat(batch_data[None, :], key.shape[0], axis=0),
              labels=jnp.repeat(
                  batch_labels[None, :], key.shape[0], axis=0
              ),
          )
  
      self._reset_fn = jax.jit(reset_fn)
  
      def get_sequence_lengths(sequence_batch):
          # sequence_batch.shape = (batch_size, seq_length, vocab_size)
          eos_row = sequence_batch[:, :, char_table.eos_id]
          eos_idx = jnp.argmax(eos_row, axis=-1)
          return jnp.where(
              eos_row[jnp.arange(eos_row.shape[0]), eos_idx],
              eos_idx + 1,
              sequence_batch.shape[1],
          )
  
      def mask_sequences(sequence_batch, lengths):
          return sequence_batch * (
              lengths[:, np.newaxis]
              > np.arange(sequence_batch.shape[1])[np.newaxis]
          )
  
      def cross_entropy_loss(logits, labels, lengths):
          xe = jnp.sum(jax.nn.log_softmax(logits) * labels, axis=-1)
          return -jnp.mean(mask_sequences(xe, lengths))
  
      def step_fn(state, action):
          labels = state.labels[:, 1:]
          lengths = get_sequence_lengths(labels)
          if test:
              token_acc = jnp.argmax(action, -1) == jnp.argmax(
                  labels, -1
              )
              sequence_acc = (
                  jnp.sum(mask_sequences(token_acc, lengths), axis=-1)
                  == lengths
              )
              reward = jnp.mean(sequence_acc)
          else:
              reward = -cross_entropy_loss(action, labels, lengths)
          done = jnp.ones((), dtype=jnp.int32)
          return state, reward, done
  
      self._step_fn = jax.jit(jax.vmap(step_fn))
  
file_path: evojax/evojax/task/seq2seq.py
imports: import numpy
  from typing import Tuple
  import jax
  import jax.numpy
  from jax import random
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
label_desc: Use image node_id 1 to create a new Seq2seqTask object from inherited base classes: VectorizedTask with example: obj = Seq2seqTask(batch_size, max_len_query_digit, test)