    def update_base_tensor(
        self, new_tensor: Tensor, stride_dim: int, stride_dim_offset: int
    ) -> None:
        """
        Updates the TensorAccessor with a new base tensor.
        This API is useful to handle ops with a stride dim, e.g. split, cat.
        It can also be used by slice if slice is only operated on one dim.
        """

        assert (
            self.stride_dim is None
        ), "Tensor accessor cannot be updated once stride_dim is set!"

        original_shapes = (
            self.actual_shapes
            if self.actual_shapes is not None
            else self.original_shapes
        )
        self.actual_shapes = new_tensor._attrs["shape"]
        self.stride_dim = stride_dim
        assert len(self.actual_shapes) == len(original_shapes), (
            f"Original tensor and new tensor must have the same number of dims! "
            f"Original tensor shape: {original_shapes}, new tensor shape: {self.actual_shapes}"
        )
        assert (
            len(self.actual_shapes) > stride_dim
        ), f"stride_dim {stride_dim} must be less than #dims {len(self.actual_shapes)}!"
        assert isinstance(
            original_shapes[stride_dim], IntImm
        ), "Stride dim can't be dynamic!"
        assert isinstance(
            self.actual_shapes[stride_dim], IntImm
        ), "Stride dim can't be dynamic!"

        self.original_total_elements_from_stride_dim = original_shapes[
            stride_dim
        ].value()
        self.actual_total_elements_from_stride_dim = self.actual_shapes[
            stride_dim
        ].value()
        self.offset = stride_dim_offset

        for original_shape, actual_shape in zip(
            original_shapes[stride_dim + 1 :], self.actual_shapes[stride_dim + 1 :]
        ):
            assert isinstance(
                original_shape, IntImm
            ), "Dims after the stride dim must have static shapes! Shapes: {}".format(
                original_shapes
            )
            assert isinstance(
                actual_shape, IntImm
            ), "Dims after the stride dim must have static shapes! Shapes: {}".format(
                self.actual_shapes
            )
            assert (
                original_shape._attrs["values"] == actual_shape._attrs["values"]
            ), "original shapes {} and actual shapes {} after the stride dim must be equal! ".format(
                original_shapes, self.actual_shapes
            )
            value = actual_shape._attrs["values"][0]
            self.original_total_elements_from_stride_dim *= value
            self.actual_total_elements_from_stride_dim *= value
            self.offset *= value

        if (
            stride_dim > 0
            and self.actual_total_elements_from_stride_dim
            > self.original_total_elements_from_stride_dim
        ):
            self.is_contiguous = False

        if (
            self.actual_total_elements_from_stride_dim
            > self.original_total_elements_from_stride_dim
        ):
            self.is_from_strided_tensor = True
Updates the TensorAccessor with a new base tensor.
This API is useful to handle ops with a stride dim, e.g. split, cat.
It can also be used by slice if slice is only operated on one dim.
python/aitemplate/compiler/tensor_accessor.py