def get_original_index(self, idx: int) -> int:
        """
        Return the original index of the input at idx in the current "inputs" list.

        Parameters
        ----------
        idx : int
            the index of an input based on the current "inputs"

        Returns
        -------
        int
            the index of this input in the "original_inputs"
        """
        num_original_inputs = len(self._attrs["original_inputs"])
        orig_idx = None
        # track the index for the "inputs" list
        curr_idx = 0
        for i in range(num_original_inputs):
            # We don't increase curr_idx if this input is removed
            if not self._attrs["input_masks"][i]:
                continue
            # We found the original index
            if curr_idx == idx:
                orig_idx = i
                break
            curr_idx += 1
        assert orig_idx is not None, f"Expected orig_idx to be non-None for idx {idx}"
        return orig_idx

        Return the original index of the input at idx in the current "inputs" list.

        Parameters
        ----------
        idx : int
            the index of an input based on the current "inputs"

        Returns
        -------
        int
            the index of this input in the "original_inputs"
        
python/aitemplate/compiler/ops/tensor/concatenate.py