[fx_importer] Add support for 0D tensors #401
Conversation
Adds an escape hatch from creating a DenseResourceElementsAttr for single value tensors into DenseElementsAttr.
| blob_name, | ||
| tensor_type, | ||
| ) | ||
| # DenseResourceElementsAttr creation doesnt support rank 0 tensors, so we use DenseElementsAttr instead. |
There was a problem hiding this comment.
The comment should start with: one element constants are more optimizable as splat DenseElementsAttr. DenseResourceElementsAttr does not support splats, so don't use it for that case. In addition, at the time of writing, it has bugs with handling 0d tensors.
| # DenseResourceElementsAttr creation doesnt support rank 0 tensors, so we use DenseElementsAttr instead. | ||
| if np_tensor.size == 1: | ||
| elements_attr = DenseElementsAttr.get( | ||
| type=TORCH_DTYPE_TO_MLIR_TYPE[tensor.dtype](), array=np_tensor |
There was a problem hiding this comment.
Can you also pass the shape= param? Keeps it from making dodgy local decisions and is marginally more robust.
| # DenseResourceElementsAttr creation doesnt support rank 0 tensors, so we use DenseElementsAttr instead. | ||
| if np_tensor.size == 1: | ||
| elements_attr = DenseElementsAttr.get( | ||
| type=TORCH_DTYPE_TO_MLIR_TYPE[tensor.dtype](), array=np_tensor |
There was a problem hiding this comment.
Wrap in try...except KeyError and raise a descriptive error message.
|
I'll need to sync this to torch-mlir once this lands. We're halfway through using that one as sot so need to port patches manually until it resolves. Can you leave the issue open and assign to me to port so I don't forget? |
| elements_attr = DenseElementsAttr.get( | ||
| type=TORCH_DTYPE_TO_MLIR_TYPE[tensor.dtype](), array=np_tensor | ||
| type=element_type, array=np_tensor, shape=[] |
There was a problem hiding this comment.
Shape param is wrong for 1d, 1 element. Get it from np_tensor.shape
There was a problem hiding this comment.
oof, sounds like a missing test case
There was a problem hiding this comment.
Yeah, add tests for shape [1].
| @@ -928,10 +928,15 @@ def _make_vtensor_literal_op( | |||
| # buffer is via the indirection: Tensor -> list -> numpy array. This allows us to create a vtensor literal as | |||
| # desired, but also limits which data types we can support in this function (see TORCH_DTYPE_TO_NPY_TYPE above) | |||
| np_tensor = np.array(tensor.tolist()).astype(npy_dtype) | |||
| # DenseResourceElementsAttr creation doesnt support rank 0 tensors, so we use DenseElementsAttr instead. | |||
| # one element constants are more optimizable as splat DenseElementsAttr. DenseResourceElementsAttr does not support splats, so don't use it for that case. In addition, at the time of writing, it has bugs with handling 0d tensors. | |||
There was a problem hiding this comment.
Capitalize first letter.
Adds an escape hatch from creating a DenseResourceElementsAttr for single value tensors into DenseElementsAttr.
Addresses #398