I'm trying to work with new FixedShapeTensors from Pyarrow. I'd like to use tensors to store covariance matrices, so I'd like to get the diagonal elements of my (6x6) matrix.
I'm working through the most naive version of this, which is to hardcode the selections of the (0,0), (1, 1), ... (6, 6) elements. I do that with "pc.list_element(pc.list_element(my_array, 0), 0)". When I try that, I get an error that I think might expose a weakness in the interactions between pyarrow.compute and extension types: pyarrow.lib.ArrowNotImplementedError: Function 'list_element' has no kernel matching input types (extension<arrow.fixed_shape_tensor>, int64) It's no surprise that there is no matching kernel for an extension type, although I'd sort of expect that pyarrow would inspect the *storage* of the extension array type. Is there a better way for me to approach this?
