How to calculate std and mean in RAy dataset

1

The following is throwing error while calculating mean and std


ray.data.from_items([
    {"A": np.array([i,i+1], dtype=np.float32), "B": i**2}
    for i in range(100)
]).mean(["A", "B"])

Giving error
_compute.pyx", line 385, in pyarrow._compute.Function.call
File “pyarrow/error.pxi”, line 154, in pyarrow.lib.pyarrow_internal_check_status
File “pyarrow/error.pxi”, line 91, in pyarrow.lib.check_status
pyarrow.lib.ArrowNotImplementedError: Function ‘sum’ has no kernel matching input types (extension<ray.data.arrow_tensor>)