def _visit_axis_update_op(self, op):
        """
        Check:
        - at least one of the non-constant inputs to this op is of type _LazyTransposeHypotheticalValue
        - for all non-constant inputs, that are of type _LazyTransposeHypotheticalValue, they
        have the same perm value.
        These checks are common for all "axis update" ops.
        """
        input_vars = self._get_input_vars(op, only_nonconst_vars=True)
        perm = None
        num_lazy_input_vars = 0
        for var in input_vars:
            hypothetical_value = self.var_to_hypothetical_value[var]
            if isinstance(hypothetical_value, _LazyTransposeHypotheticalValue):
                num_lazy_input_vars += 1
                if perm is None:
                    perm = hypothetical_value.perm
                elif perm != hypothetical_value.perm:
                    self._visit_materialize_op(op)
                    return

        if num_lazy_input_vars == 0:
            self._visit_materialize_op(op)
            return

        # checks specific to the op type
        op_cls = self._AXIS_UPDATE_OPS.get(op.op_type, None)
        if op_cls is None:
            raise ValueError("Transform class for op of type '{}' not found".format(op.op_type))

        if not op_cls(
            **{
                "op": op,
                "transpose_axes": perm,
                "var_to_hypothetical_value_dict": self.var_to_hypothetical_value,
            }
        ).can_transpose_pass():
            self._visit_materialize_op(op)
            return

        # add this op to the dictionary "transpose_op_to_axis_update_ops"
        # and update self.var_to_hypothetical_value[op.outputs[0]]
        all_lazy_transpose_ops = set()
        wrapped_hypothetical_value = None
        for var in input_vars:
            input_hypothetical_value = self.var_to_hypothetical_value[var]
            if isinstance(input_hypothetical_value, _LazyTransposeHypotheticalValue):
                all_lazy_transpose_ops.update(input_hypothetical_value.transpose_ops)
                wrapped_hypothetical_value = input_hypothetical_value.wrapped_hypothetical_value

        for transpose_op in all_lazy_transpose_ops:
            self.transpose_op_to_axis_update_ops[transpose_op].append(op)

        for output in op.outputs:
            self.var_to_hypothetical_value[output] = _LazyTransposeHypotheticalValue(
                wrapped_hypothetical_value,
                all_lazy_transpose_ops,
                perm,
            )

        Check:
        - at least one of the non-constant inputs to this op is of type _LazyTransposeHypotheticalValue
        - for all non-constant inputs, that are of type _LazyTransposeHypotheticalValue, they
        have the same perm value.
        These checks are common for all "axis update" ops.
        
coremltools/converters/mil/mil/passes/defs/optimize_repeat_ops.py