def get_io_shape(self, io_name: str) -> Optional[list[Union[int, str]]]:
        """Get the shape of an input/output.

        :param io_name: name of the input/output.
        :return: shape of the input/output.
        """
        if self.is_initializer(io_name):
            return list(self.get_initializer_proto(io_name).dims)

        tensor_type = self._get_vi_tensor_type(io_name)
        return (
            [dim.dim_param if dim.dim_param else dim.dim_value for dim in tensor_type.shape.dim]
            if tensor_type
            else None
        )
Get the shape of an input/output.

        :param io_name: name of the input/output.
        :return: shape of the input/output.
        
olive/passes/onnx/onnx_dag.py