Lunderberg commented on code in PR #16563:
URL: https://github.com/apache/tvm/pull/16563#discussion_r1491708293
##########
python/tvm/relax/expr.py:
##########
@@ -244,6 +244,192 @@ def __getitem__(self, index: int) -> "ExprWithOp":
raise IndexError from err
raise
+ def _check_for_tensor_struct_info(self):
+ """Raise an error if this is something other than a Tensor
+
+ Used for early checks in `expr.dtype` and `expr.shape`
+ accessors. While invalid usage would cause errors to be
+ raised durin shape inference, an earlier check makes it easier
+ to find the invalid usage.
+ """
+ if self.struct_info_ is None:
+ return
+
+ if not isinstance(self.struct_info_, tvm.relax.TensorStructInfo):
+ raise TypeError(
+ f"Runtime unpacking of DLDataType is only implemented for
tensors, "
+ f"but was applied to object {self} of type {type(self)}."
+ )
+
+ @property
+ def dtype(self) -> "_DLTensorDTypeProxy":
+ """Returns a proxy object for accessing DLTensor::dtype"""
+ self._check_for_tensor_struct_info()
+ return _DLTensorDTypeProxy(self)
+
+ @property
+ def ndim(self) -> "Expr":
+ """Returns the runtime value of DLTensor::ndim"""
+ self._check_for_tensor_struct_info()
+ op = tvm.ir.Op.get("relax.tensor_ndim")
+ return tvm.relax.Call(op, [self])
+
+ @property
+ def shape(self) -> "_DLTensorShapeProxy":
+ """Returns a proxy object for accessing DLTensor::shape"""
+ self._check_for_tensor_struct_info()
+ return _DLTensorShapeProxy(self)
+
+
+class _DLTensorDTypeProxy(tvm.runtime.ObjectGeneric):
+ """A proxy object for unpacking DLDatatype from DLTensor
+
+ Exposes accessors for `DLDataType` fields `type_code`, `lanes`,
+ and `bits` within a `DLTensor::dtype`. Accessing these fields
Review Comment:
Good point. At the moment, I stuck with the values that have a direct
presence elsewhere in Relax, but I agree that it would be good to be able to
extract any of the `DLTensor*` fields.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
