Class ShapeProxy

python/tvm/relax/script/parser/entry.py:407–437

class ShapeProxy(StructInfoProxy):
    values: list[PrimExpr] | None
    ndim: int
    """The type of shape values.

    Parameters
    ----------
    values : Optional[List[PrimExpr]]
       The symbolic shape values if known.

    ndim : Optional[int]
       The size of the shape.
    """

    def __init__(
        self,
        values: list[PrimExpr] | None = None,
        ndim: int = -1,
    ) -> None:
        self.values = values
        self.ndim = ndim

    def get_symbolic_vars(self) -> set[str]:
        if self.values is None:
            return set()
        else:
            return {v for v in self.values if isinstance(v, str) and v.isidentifier()}

    def as_struct_info(self, dict_globals: dict[str, Any] | None = None) -> ShapeStructInfo:
        values = [_eval_shape(v, dict_globals) for v in self.values] if self.values else None
        return ShapeStructInfo(values, self.ndim)


def Shape(values: list[PrimExpr] | None = None, ndim: int = -1) -> ShapeProxy: