def AudioSpectrogram(context, node):
    """
    input shape: (Tin, channels)
    attributes: stride (int), window_size (int), magnitude_squared (bool)

    output shape : (channels, Tout, fout)
    where,
    Tout = floor((Tin - window_size)/stride + 1)
    fout = N / 2 + 1
    where N = next_power_of_2(window_size) = 2 ^ ceil(log2(window_size))

    reference:
    https://github.com/tensorflow/tensorflow/blob/dec8e0b11f4f87693b67e125e67dfbc68d26c205/tensorflow/core/kernels/spectrogram_op.cc
    """

    x = context[node.inputs[0]] # (Tin, channels)
    if x.rank != 2:
        raise NotImplementedError("AudioSpectrogram op: rank of the input must be 2")

    if "magnitude_squared" not in node.attr:
        raise ValueError("AudioSpectrogram op: missing attribute: 'magnitude_squared'")
    if "stride" not in node.attr:
        raise ValueError("AudioSpectrogram op: missing attribute: 'stride'")
    if "window_size" not in node.attr:
        raise ValueError("AudioSpectrogram op: missing attribute: 'window_size'")

    magnitude_squared = node.attr["magnitude_squared"]
    stride = node.attr["stride"]
    window_size = node.attr["window_size"]

    N = 2 ** _np.ceil(_np.log2(window_size))
    N = N.astype(_np.int32)
    fout = N / 2 + 1
    fout = fout.astype(_np.int32)

    # construct constant for hann window tensor, of shape (window_size,)
    h = _np.arange(window_size) * ((2 * _np.pi) / window_size)
    h = 0.5 - 0.5 * _np.cos(h)

    # construct the constant DFT matrices
    k = _np.arange(fout).reshape(1, fout) # (1, fout)
    n = _np.arange(N).reshape(N, 1)  # (N, 1)
    kn = _np.matmul(n, k) * (2 * _np.pi / N) # (N, fout)
    Re_DFT_matrix_const = _np.cos(kn)  # (N, fout)
    Im_DFT_matrix_const = -_np.sin(kn) # (N, fout)

    # transpose input
    x = mb.transpose(x=x, perm=[1,0]) # (channels, Tin)
    # extract slices from the input
    x = mb.sliding_windows(x=x, axis=1, size=window_size, stride=stride) # (channels, Tout, window_size)
    # multiply with hann window
    x = mb.mul(x=x, y=h)
    # pad the last dimension to size N
    x = mb.pad(x=x, pad=[0,0,0,0,0,N - window_size], mode="constant", constant_val=0.0) # (channels, Tout, N)
    # multiply by DFT matrices
    re = mb.matmul(x=x, y=Re_DFT_matrix_const) # (channels, Tout, fout)
    im = mb.matmul(x=x, y=Im_DFT_matrix_const)  # (channels, Tout, fout)

    # compute spectrogram
    re = mb.mul(x=re, y=re)
    im = mb.mul(x=im, y=im)
    if not magnitude_squared:
        y = mb.add(x=re, y=im)
        y = mb.sqrt(x=y, name=node.name)
    else:
        y = mb.add(x=re, y=im, name=node.name)
    context.add(node.name, y)

    input shape: (Tin, channels)
    attributes: stride (int), window_size (int), magnitude_squared (bool)

    output shape : (channels, Tout, fout)
    where,
    Tout = floor((Tin - window_size)/stride + 1)
    fout = N / 2 + 1
    where N = next_power_of_2(window_size) = 2 ^ ceil(log2(window_size))

    reference:
    https://github.com/tensorflow/tensorflow/blob/dec8e0b11f4f87693b67e125e67dfbc68d26c205/tensorflow/core/kernels/spectrogram_op.cc
    
coremltools/converters/mil/frontend/tensorflow/ops.py