call method

Implementation

({MlxArray features, List<int> lengths}) call(
  MlxArray mel,
  List<int> lengths,
) {
  var nextLengths = List<int>.from(lengths);
  final input = mel.expandDims(1); // [B, 1, T, F]
  final layoutInput = input.transposeAxes(<int>[0, 2, 3, 1]); // [B, T, F, C]
  input.close();

  final conv0 = _conv0(
    layoutInput,
    stride: <int>[2, 2],
    padding: <int>[1, 1],
  );
  layoutInput.close();
  final act0 = parakeetRelu(conv0);
  conv0.close();
  nextLengths = nextLengths.map((len) => _convLen(len, 3, 2, 1)).toList();

  final conv2 = _conv2(act0, stride: <int>[2, 2], padding: <int>[1, 1]);
  act0.close();
  final conv3 = _conv3(conv2, stride: <int>[1, 1], padding: <int>[0, 0]);
  conv2.close();
  final act1 = parakeetRelu(conv3);
  conv3.close();
  nextLengths = nextLengths.map((len) => _convLen(len, 3, 2, 1)).toList();

  final conv5 = _conv5(act1, stride: <int>[2, 2], padding: <int>[1, 1]);
  act1.close();
  final conv6 = _conv6(conv5, stride: <int>[1, 1], padding: <int>[0, 0]);
  conv5.close();
  final act2 = parakeetRelu(conv6);
  conv6.close();
  nextLengths = nextLengths.map((len) => _convLen(len, 3, 2, 1)).toList();

  final batch = act2.shape[0];
  final time = act2.shape[1];
  final freq = act2.shape[2];
  final channels = act2.shape[3];
  final flattened = act2.transposeAxes(<int>[0, 1, 3, 2]).reshape(<int>[
    batch,
    time,
    channels * freq,
  ]);
  act2.close();
  final projected = _out(
    flattened.reshape(<int>[batch * time, channels * freq]),
  );
  flattened.close();
  final features = projected.reshape(<int>[
    batch,
    time,
    bundle.manifest.encoderHidden,
  ]);
  projected.close();
  return (features: features, lengths: nextLengths);
}