class_docstr: Predict melspectrogram from aligned phonemes
class_docstr_tok: ["Predict","melspectrogram","from","aligned","phonemes"]
code: def __call__(self, inputs: AcousticInput):
      x = self.encoder(inputs.phonemes, inputs.lengths)
      x = self.upsample(x, inputs.durations, inputs.mels.shape[1])
      mels = self.prenet(inputs.mels)
      x = jnp.concatenate((x, mels), axis=-1)
      B, L, _ = x.shape
      hx = self.decoder.initial_state(B)
  
      def zoneout_decoder(inputs, prev_state):
          x, mask = inputs
          x, state = self.decoder(x, prev_state)
          state = jax.tree_multimap(
              lambda m, s1, s2: s1 * m + s2 * (1 - m),
              mask,
              prev_state,
              state,
          )
          return x, state
  
      mask = jax.tree_map(
          lambda x: jax.random.bernoulli(
              hk.next_rng_key(), 0.1, (B, L, x.shape[-1])
          ),
          hx,
      )
      x, _ = hk.dynamic_unroll(
          zoneout_decoder, (x, mask), hx, time_major=False
      )
      x = self.projection(x)
      residual = self.postnet(x)
      return x, x + residual
  
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 6 for calling the AcousticModel obj's underlying member method code with example usage: obj.__call__(inputs) and returns: x,  unknown, x,  state