code: def loss_fn(
      params, aux, rng, inputs: AcousticInput, is_training=True
  ):
      melfilter = MelFilter(
          FLAGS.sample_rate,
          FLAGS.n_fft,
          FLAGS.mel_dim,
          FLAGS.fmin,
          FLAGS.fmax,
      )
      mels = melfilter(inputs.wavs.astype(jnp.float32) / (2**15))
      B, L, D = mels.shape
      go_frame = jnp.zeros((B, 1, D), dtype=jnp.float32)
      inp_mels = jnp.concatenate((go_frame, mels[:, :-1, :]), axis=1)
      n_frames = (
          inputs.durations * FLAGS.sample_rate / (FLAGS.n_fft // 4)
      )
      inputs = inputs._replace(mels=inp_mels, durations=n_frames)
      model = net if is_training else val_net
      (mel1_hat, mel2_hat), new_aux = model.apply(
          params, aux, rng, inputs
      )
      loss1 = (
          jnp.square(mel1_hat - mels) + jnp.square(mel2_hat - mels)
      ) / 2
      loss2 = (jnp.abs(mel1_hat - mels) + jnp.abs(mel2_hat - mels)) / 2
      loss = jnp.mean((loss1 + loss2) / 2, axis=-1)
      num_frames = (inputs.wav_lengths // (FLAGS.n_fft // 4))[:, None]
      mask = jnp.arange(0, L)[None, :] < num_frames
      loss = jnp.sum(loss * mask) / jnp.sum(mask)
      return (
          (loss, new_aux)
          if is_training
          else (loss, new_aux, mel2_hat, mels)
      )
  
file_path: underthesea/underthesea/pipeline/say/viettts_/nat/acoustic_trainer.py
imports: import pickle
  from functools import partial
  from typing import Deque
  import haiku
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import optax
  from tqdm.auto import tqdm
  from vietTTS.nat.config import AcousticInput
  from .config import FLAGS, AcousticInput
  from .data_loader import load_textgrid_wav
  from .dsp import MelFilter
  from .model import AcousticModel
  from .utils import print_flags
label_desc: Use image node_id 3 for calling a global function with example usage: loss_fn(params, aux, rng, inputs, is_training) and returns: unknown