debugStages method

Implementation

Map<String, MlxArray> debugStages(MlxArray mel) {
  final input = mel.expandDims(1);
  final layoutInput = input.transposeAxes(<int>[0, 2, 3, 1]);
  input.close();
  final conv0 = _conv0(
    layoutInput,
    stride: <int>[2, 2],
    padding: <int>[1, 1],
  );
  layoutInput.close();
  final act0 = parakeetRelu(conv0);
  final conv2 = _conv2(act0, stride: <int>[2, 2], padding: <int>[1, 1]);
  final conv3 = _conv3(conv2, stride: <int>[1, 1], padding: <int>[0, 0]);
  final act1 = parakeetRelu(conv3);
  final conv5 = _conv5(act1, stride: <int>[2, 2], padding: <int>[1, 1]);
  final conv6 = _conv6(conv5, stride: <int>[1, 1], padding: <int>[0, 0]);
  final act2 = parakeetRelu(conv6);
  return <String, MlxArray>{
    'conv0': conv0,
    'act0': act0,
    'conv2': conv2,
    'conv3': conv3,
    'act1': act1,
    'conv5': conv5,
    'conv6': conv6,
    'act2': act2,
  };
}