ekalda commented on a change in pull request #9457:
URL: https://github.com/apache/tvm/pull/9457#discussion_r750115698
##########
File path: python/tvm/relay/backend/contrib/ethosu/te/dma.py
##########
@@ -67,63 +67,72 @@ def _pad(*indices):
return _pad
-def read_compute(tensor: te.Tensor, layout: str, zero_point: int, scale:
float) -> te.Tensor:
+def read_compute(tensor: te.Tensor, zero_point: int, scale: float, layout: str
= None) -> te.Tensor:
Review comment:
Done
##########
File path: python/tvm/relay/backend/contrib/ethosu/te/dma.py
##########
@@ -67,63 +67,72 @@ def _pad(*indices):
return _pad
-def read_compute(tensor: te.Tensor, layout: str, zero_point: int, scale:
float) -> te.Tensor:
+def read_compute(tensor: te.Tensor, zero_point: int, scale: float, layout: str
= None) -> te.Tensor:
"""A tensor expression which represents a read.
Parameters
----------
tensor : te.Tensor
The tensor to read.
- layout : str
- The layout of the tensor, either NHWC or NHCWB16.
zero_point : int
The zero point of the tensor.
scale : float
The scale of the tensor.
+ layout : Optional[str]
+ The layout of the tensor, either NHWC or NHCWB16.
Returns
-------
te.Tensor
The tensor having been read.
"""
- assert layout in {"NHWC", "NHCWB16"}
read_attrs = {
"op": "ethosu_read",
- "layout": layout,
"zero_point": zero_point,
"scale": scale,
}
+
+ if layout:
+ assert layout in {"NHWC", "NHCWB16"}
+ read_attrs["layout"] = layout
+
return te.compute(tensor.shape, lambda *i: tensor(*i), name="ethosu_read",
attrs=read_attrs)
-def write_compute(tensor: te.Tensor, layout: str, zero_point: int, scale:
float) -> te.Tensor:
+def write_compute(
+ tensor: te.Tensor, zero_point: int, scale: float, layout: str = None
Review comment:
Done
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
