    def _try_gen_dim_mapping(self):
        """
        Updates self._dim_mapping based on self.original_shapes and self.actual_shapes.
        Check comments for self._dim_mapping for more info.
        """

        # Set self._dim_mapping = None.
        self._dim_mapping = None
        dim_mapping = []

        original_value = None
        actual_value = None
        original_idx = 0
        actual_idx = 0
        prev_original_idx = 0
        prev_actual_idx = 0

        # Add a dummy entry at the end of original_shapes / actual_shapes to
        # make boundary checks easier.
        INT_MAX = int("0x" + "F" * 16, 16)
        original_shapes = list(self.original_shapes) + [IntImm(INT_MAX)]
        actual_shapes = list(self.actual_shapes) + [IntImm(INT_MAX)]

        while original_idx < len(original_shapes) and actual_idx < len(actual_shapes):
            original_d = original_shapes[original_idx]
            actual_d = actual_shapes[actual_idx]

            if original_d == IntImm(1) and original_value is None:
                # Create separate groups for "1"s.
                original_idx += 1
                dim_mapping.append((list(range(prev_original_idx, original_idx)), []))
                prev_original_idx = original_idx
                continue

            if actual_d == IntImm(1) and actual_value is None:
                # Create separate groups for "1"s.
                actual_idx += 1
                dim_mapping.append(([], list(range(prev_actual_idx, actual_idx))))
                prev_actual_idx = actual_idx
                continue

            if not isinstance(original_d, IntImm) or not isinstance(actual_d, IntImm):
                # If there are dynamic dims, check whether these two dynamic dims are the same.
                # Only allow one dynamic dim per group.
                if (
                    original_d != actual_d
                    or original_value is not None
                    or actual_value is not None
                ):
                    return
                else:
                    original_value = original_d
                    actual_value = actual_d
                    original_idx += 1
                    actual_idx += 1
            elif original_value is None or actual_value is None:
                # Assign values to original_value / actual_value first.
                if original_value is None:
                    original_value = original_d.value()
                    original_idx += 1
                if actual_value is None:
                    actual_value = actual_d.value()
                    actual_idx += 1
            else:
                # Choose to advance original_idx or actual_idx based on
                # original_value / actual_value comparisons.
                if original_value < actual_value:
                    original_value *= original_d.value()
                    original_idx += 1
                elif original_value > actual_value:
                    actual_value *= actual_d.value()
                    actual_idx += 1
                else:
                    raise AssertionError("This branch should never be reached.")

            if original_value == actual_value:
                dim_mapping.append(
                    (
                        list(range(prev_original_idx, original_idx)),
                        list(range(prev_actual_idx, actual_idx)),
                    )
                )
                prev_original_idx = original_idx
                prev_actual_idx = actual_idx
                original_value = None
                actual_value = None

        if (
            original_value is not None
            or actual_value is not None
            or original_idx != len(original_shapes)
            or actual_idx != len(actual_shapes)
        ):
            _LOGGER.debug(f"tail processing failed, dim_mapping: {dim_mapping}")
            return

        # Remove the last dummy group.
        dim_mapping = dim_mapping[:-1]

        # Assign new dim_mapping to self._dim_mapping.
        self._dim_mapping = dim_mapping
        _LOGGER.debug(f"generate dim_mapping: {dim_mapping}")
Updates self._dim_mapping based on self.original_shapes and self.actual_shapes.
Check comments for self._dim_mapping for more info.
python/aitemplate/compiler/tensor_accessor.py