    def _replace_jagged_int_var(shape: List[IntVar]):
        """
        If shape[0] is a JaggedIntVar, replace it with
        the corresponding maximum dense shape.
        """
        if shape and isinstance(shape[0], JaggedIntVar):
            return shape[0].get_max_dense_shape() + shape[1:]
        return shape
If shape[0] is a JaggedIntVar, replace it with
the corresponding maximum dense shape.
python/aitemplate/compiler/transform/remove_no_ops.py