How to use the nemo.core.neural_types.AxisType function in NEMO

def create_ports():
        input_ports = {
            "encoder_output": NeuralType(
                {0: AxisType(BatchTag),
                 1: AxisType(EncodedRepresentationTag),
                 2: AxisType(ProcessedTimeTag)})}
        output_ports = {
            "output": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag),
                2: AxisType(ChannelTag)
            })}
        return input_ports, output_ports

}),
            # 'tgt_lens': NeuralType({
            #     0: AxisType(BatchTag),
            #     1: AxisType(ChannelTag)
            # }),
            'targets': NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(ChannelTag),  # the number of slots
                2: AxisType(TimeTag)
            }),

        }
        output_ports = {
            'point_outputs': NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag),
                2: AxisType(ChannelTag),
                3: AxisType(ChannelTag)
            }),
            'gate_outputs': NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag),
                2: AxisType(ChannelTag)
            })
        }
        return input_ports, output_ports

input_ports = {
            'targets': NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag)
            }),
            'encoder_outputs': NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag),
                2: AxisType(ChannelTag)
            }, optional=True)
        }
        output_ports = {
            'log_probs': NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag),
                2: AxisType(ChannelTag)
            }),
            'attention_weights': NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag),
                2: AxisType(TimeTag)
            }, optional=True)
        }
        return input_ports, output_ports

"in_y": NeuralType({0: AxisType(BatchTag),
                                1: AxisType(TimeTag)}),

            "in_y_len": NeuralType({0: AxisType(BatchTag)})
        }

        output_ports = {
            "out_x": NeuralType({0: AxisType(BatchTag),
                                 1: AxisType(SpectrogramSignalTag),
                                 2: AxisType(TimeTag)}),

            "out_x_len": NeuralType({0: AxisType(BatchTag)}),

            "out_y": NeuralType({0: AxisType(BatchTag),
                                 1: AxisType(TimeTag)}),

            "out_y_len": NeuralType({0: AxisType(BatchTag)})
        }
        return input_ports, output_ports

def create_ports():
        input_ports = {
            "audio_signal": NeuralType({0: AxisType(BatchTag),
                                        1: AxisType(SpectrogramSignalTag),
                                        2: AxisType(ProcessedTimeTag)}),
            "length": NeuralType({0: AxisType(BatchTag)})
        }

        output_ports = {
            "outputs": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(EncodedRepresentationTag),
                2: AxisType(ProcessedTimeTag),
            }),

            "encoded_lengths": NeuralType({0: AxisType(BatchTag)})
        }
        return input_ports, output_ports

def create_ports():
        input_ports = {
            "preds": NeuralType({
                0: AxisType(RegressionTag)
            }),
            "labels": NeuralType({
                0: AxisType(RegressionTag)
            })
        }

        output_ports = {
            "loss": NeuralType(None)
        }
        return input_ports, output_ports

def create_ports():
        input_ports = {
            "logits": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag),
                2: AxisType(ChannelTag),
                3: AxisType(ChannelTag)
            }),
            "targets": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(ChannelTag),
                2: AxisType(TimeTag)
            }),
            "mask": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(ChannelTag)
            }),
        }

        output_ports = {"loss": NeuralType(None)}
        return input_ports, output_ports

def create_ports():
        input_ports = {
            "log_probs": NeuralType({0: AxisType(BatchTag),
                                     1: AxisType(TimeTag),
                                     2: AxisType(ChannelTag)})
        }

        output_ports = {
            "predictions": NeuralType({
                0: AxisType(BatchTag),
                1: AxisType(TimeTag)
            })
        }
        return input_ports, output_ports