class_docstr: Encode phonemes/text to vector
class_docstr_tok: ["Encode","phonemes\/text","to","vector"]
code: def __call__(self, x, lengths):
      x = self.embed(x)
      x = jax.nn.relu(
          self.bn1(self.conv1(x), is_training=self.is_training)
      )
      if self.is_training:
          x = hk.dropout(hk.next_rng_key(), self.dropout_rate, x)
      x = jax.nn.relu(
          self.bn2(self.conv2(x), is_training=self.is_training)
      )
      if self.is_training:
          x = hk.dropout(hk.next_rng_key(), self.dropout_rate, x)
      x = jax.nn.relu(
          self.bn3(self.conv3(x), is_training=self.is_training)
      )
      if self.is_training:
          x = hk.dropout(hk.next_rng_key(), self.dropout_rate, x)
      B, L, _ = x.shape
      mask = jnp.arange(0, L)[None, :] >= (lengths[:, None] - 1)
      h0c0_fwd = self.lstm_fwd.initial_state(B)
      new_hx_fwd, _ = hk.dynamic_unroll(
          self.lstm_fwd, x, h0c0_fwd, time_major=False
      )
      x_bwd, mask_bwd = jax.tree_map(
          lambda x: jnp.flip(x, axis=1), (x, mask)
      )
      h0c0_bwd = self.lstm_bwd.initial_state(B)
      new_hx_bwd, _ = hk.dynamic_unroll(
          self.lstm_bwd, (x_bwd, mask_bwd), h0c0_bwd, time_major=False
      )
      x = jnp.concatenate(
          (new_hx_fwd, jnp.flip(new_hx_bwd, axis=1)), axis=-1
      )
      return x
  
imports: import haiku
  import jax
  import jax.numpy
  from jax.numpy import ndarray
  from .config import FLAGS, AcousticInput, DurationInput
label_desc: Use image node_id 2 for calling the TokenEncoder obj's underlying member method code with example usage: obj.__call__(x, lengths) and returns: x