class VADRNNJITMerge(Module):
  __parameters__ = []
  __buffers__ = []
  training : bool
  _is_full_backward_hook : Optional[bool]
  _h : Tensor
  _c : Tensor
  _last_sr : int
  _last_batch_size : int
  sample_rates : List[int]
  _model : __torch__.vad.model.vad_annotator.___torch_mangle_28.VADRNNJIT
  _model_8k : __torch__.vad.model.vad_annotator.VADRNNJIT
  def forward(self: __torch__.vad.model.vad_annotator.VADRNNJITMerge,
    x: Tensor,
    sr: int) -> Tensor:
    _0 = uninitialized(Tensor)
    x0, sr0, = (self)._validate_input(x, sr, )
    _last_sr = self._last_sr
    if bool(_last_sr):
      _last_sr0 = self._last_sr
      _1 = torch.ne(_last_sr0, sr0)
    else:
      _1 = False
    if _1:
      _2 = (self).reset_states()
    else:
      pass
    _last_batch_size = self._last_batch_size
    if bool(_last_batch_size):
      _last_batch_size0 = self._last_batch_size
      _4 = torch.ne(_last_batch_size0, (torch.size(x0))[0])
      _3 = _4
    else:
      _3 = False
    if _3:
      _5 = (self).reset_states()
    else:
      pass
    if torch.eq(sr0, 16000):
      _model = self._model
      _h = self._h
      _c = self._c
      out0, _6, _7, = (_model).forward(x0, _h, _c, )
      self._h = _6
      self._c = _7
      out = out0
    else:
      if torch.eq(sr0, 8000):
        _model_8k = self._model_8k
        _h0 = self._h
        _c0 = self._c
        _8 = (_model_8k).forward(x0, _h0, _c0, )
        out2, _9, _10, = _8
        self._h = _9
        self._c = _10
        out1 = out2
      else:
        ops.prim.RaiseException("", "builtins.ValueError")
        out1 = _0
      out = out1
    self._last_sr = sr0
    _h1 = self._h
    self._last_batch_size = (torch.size(_h1))[1]
    return out
  def _validate_input(self: __torch__.vad.model.vad_annotator.VADRNNJITMerge,
    x: Tensor,
    sr: int) -> Tuple[Tensor, int]:
    _11 = "Too many dimensions for input audio chunk {}"
    _12 = "Supported sampling rates: {} (or multiply of 16000)"
    if torch.eq(torch.dim(x), 1):
      x1 = torch.unsqueeze(x, 0)
    else:
      x1 = x
    if torch.gt(torch.dim(x1), 2):
      ops.prim.RaiseException(torch.format(_11, torch.dim(x1)), "builtins.ValueError")
    else:
      pass
    if torch.ne(sr, 16000):
      _14 = torch.eq(torch.remainder(sr, 16000), 0)
      _13 = _14
    else:
      _13 = False
    if _13:
      step = torch.floordiv(sr, 16000)
      x3 = torch.slice(x1, 0, None, None, step)
      sr1, x2 = 16000, x3
    else:
      sr1, x2 = sr, x1
    sample_rates = self.sample_rates
    _15 = torch.__contains__(sample_rates, sr1)
    if torch.__not__(_15):
      sample_rates0 = self.sample_rates
      ops.prim.RaiseException(torch.format(_12, sample_rates0), "builtins.ValueError")
    else:
      pass
    _16 = torch.gt(torch.div(sr1, (torch.size(x2))[1]), 31.25)
    if _16:
      ops.prim.RaiseException("Input audio chunk is too short", "builtins.ValueError")
    else:
      pass
    return (x2, sr1)
  def audio_forward(self: __torch__.vad.model.vad_annotator.VADRNNJITMerge,
    x: Tensor,
    sr: int,
    num_samples: int=512) -> Tensor:
    _17 = annotate(List[Tensor], [])
    x4, sr2, = (self)._validate_input(x, sr, )
    _18 = torch.remainder((torch.size(x4))[1], num_samples)
    if bool(_18):
      _19 = torch.remainder((torch.size(x4))[1], num_samples)
      pad_num = torch.sub(num_samples, _19)
      x6 = torch.pad(x4, [0, pad_num], "constant", 0.)
      x5 = x6
    else:
      x5 = x4
    _20 = (self).reset_states()
    _21 = torch.__range_length(0, (torch.size(x5))[1], num_samples)
    for _22 in range(_21):
      i = torch.__derive_index(_22, 0, num_samples)
      wavs_batch = torch.slice(torch.slice(x5), 1, i, torch.add(i, num_samples))
      out_chunk = (self).forward(wavs_batch, sr2, )
      _23 = torch.append(_17, out_chunk)
    stacked = torch.cat(_17, 1)
    return torch.cpu(stacked)
  def reset_states(self: __torch__.vad.model.vad_annotator.VADRNNJITMerge) -> NoneType:
    self._h = torch.zeros([0])
    self._c = torch.zeros([0])
    self._last_sr = 0
    self._last_batch_size = 0
    return None
class VADDecoderRNNJIT(Module):
  __parameters__ = []
  __buffers__ = []
  training : bool
  _is_full_backward_hook : Optional[bool]
  rnn : __torch__.torch.nn.modules.rnn.LSTM
  decoder : __torch__.torch.nn.modules.container.___torch_mangle_26.Sequential
  def forward(self: __torch__.vad.model.vad_annotator.VADDecoderRNNJIT,
    x: Tensor,
    h: Tensor=CONSTANTS.c0,
    c: Tensor=CONSTANTS.c0) -> Tuple[Tensor, Tensor, Tensor]:
    if bool(torch.len(h)):
      rnn = self.rnn
      _24 = (rnn).forward__0(torch.permute(x, [0, 2, 1]), (h, c), )
      x1, _25, = _24
      h1, c1, = _25
      x0, h0, c0 = x1, h1, c1
    else:
      rnn0 = self.rnn
      _26 = (rnn0).forward__0(torch.permute(x, [0, 2, 1]), None, )
      x2, _27, = _26
      h2, c2, = _27
      x0, h0, c0 = x2, h2, c2
    x3 = torch.permute(x0, [0, 2, 1])
    decoder = self.decoder
    x4 = (decoder).forward(x3, )
    return (x4, h0, c0)
class VADRNNJIT(Module):
  __parameters__ = []
  __buffers__ = []
  training : bool
  _is_full_backward_hook : Optional[bool]
  sample_rate : int
  adaptive_normalization : __torch__.models.stt_model_blocks.___torch_mangle_0.AdaptiveAudioNormalizationNew
  feature_extractor : __torch__.vad.utils.pytorch_stft.___torch_mangle_1.STFT
  first_layer : __torch__.torch.nn.modules.container.___torch_mangle_3.Sequential
  encoder : __torch__.torch.nn.modules.container.___torch_mangle_24.Sequential
  decoder : __torch__.vad.model.vad_annotator.___torch_mangle_38.VADDecoderRNNJIT
  def forward(self: __torch__.vad.model.vad_annotator.VADRNNJIT,
    x: Tensor,
    h: Tensor=CONSTANTS.c0,
    c: Tensor=CONSTANTS.c0) -> Tuple[Tensor, Tensor, Tensor]:
    feature_extractor = self.feature_extractor
    x0 = (feature_extractor).forward(x, )
    adaptive_normalization = self.adaptive_normalization
    norm = (adaptive_normalization).forward(x0, )
    x1 = torch.cat([x0, norm], 1)
    first_layer = self.first_layer
    x2 = (first_layer).forward(x1, )
    encoder = self.encoder
    x3 = (encoder).forward(x2, )
    decoder = self.decoder
    x4, h0, c0, = (decoder).forward(x3, h, c, )
    _28 = torch.mean(torch.squeeze(x4, 1), [1])
    out = torch.unsqueeze(_28, 1)
    return (out, h0, c0)
  def audio_forward(self: __torch__.vad.model.vad_annotator.VADRNNJIT,
    x: Tensor,
    num_samples: int=512) -> Tensor:
    _29 = annotate(List[Tensor], [])
    hn = torch.zeros([0])
    cn = torch.zeros([0])
    if torch.eq(torch.dim(x), 1):
      x5 = torch.unsqueeze(x, 0)
    else:
      x5 = x
    _30 = torch.remainder((torch.size(x5))[1], num_samples)
    if bool(_30):
      _31 = torch.remainder((torch.size(x5))[1], num_samples)
      pad_num = torch.sub(num_samples, _31)
      x7 = torch.pad(x5, [0, pad_num], "constant", 0.)
      x6 = x7
    else:
      x6 = x5
    _32 = torch.__range_length(0, (torch.size(x6))[1], num_samples)
    cn0 = cn
    hn0 = hn
    for _33 in range(_32):
      i = torch.__derive_index(_33, 0, num_samples)
      wavs_batch = torch.slice(torch.slice(x6), 1, i, torch.add(i, num_samples))
      _34 = (self).forward(wavs_batch, hn0, cn0, )
      out_chunk, hn1, cn1, = _34
      _35 = torch.append(_29, out_chunk)
      cn0, hn0 = cn1, hn1
    stacked = torch.cat(_29, 1)
    return torch.cpu(stacked)
