dtype of tensors should be preserved
Author: BramVanroy
State: closed

After switching to `datasets` my model just broke. After a weekend of debugging, the issue was that my model could not handle the double that the Dataset provided, as it expected a float (but didn't give a warning, which seems a [PyTorch issue](https://discuss.pytorch.org/t/is-it-required-that-input-and-hidden-for-gru-have-the-same-dtype-float32/96221)). 

As a user I did not expect this bug. I have a `map` function that I call on the Dataset that looks like this:

```python
def preprocess(sentences: List[str]):
    token_ids = [[vocab.to_index(t) for t in s.split()] for s in sentences]

    sembeddings = stransformer.encode(sentences)
    print(sembeddings.dtype)
    return {"input_ids": token_ids, "sembedding": sembeddings}
```

Given a list of `sentences` (`List[str]`), it converts those into token_ids on the one hand (list of lists of ints; `List[List[int]]`) and into sentence embeddings on the other (Tensor of dtype `torch.float32`). That means that I actually set the column "sembedding" to a tensor that I as a user expect to be a float32.

It appears though that behind the scenes, this tensor is converted into a **list**. I did not find this documented anywhere but I might have missed it. From a user's perspective this is incredibly important though, because it means you cannot do any data_type or tensor casting yourself in a mapping function! Furthermore, this can lead to issues, as was my case. 

My model expected float32 precision, which I thought `sembedding` was because that is what `stransformer.encode` outputs. But behind the scenes this tensor is first cast to a list, and when we then set its format, as below, this column is cast not to float32 but to double precision float64.

```python
dataset.set_format(type="torch", columns=["input_ids", "sembedding"])
```

This happens because apparently there is an intermediate step of casting to a **numpy** array (?) **whose dtype creation/deduction is different from torch dtypes** (see the snippet below).  As you can see, this means that the dtype is not preserved: if I got it right, the dataset goes from torch.float32 -> list -> float64 (numpy) -> torch.float64. 

```python
import torch
import numpy as np

l = [-0.03010837361216545, -0.035979013890028, -0.016949838027358055]
torch_tensor = torch.tensor(l)
np_array = np.array(l)
np_to_torch = torch.from_numpy(np_array)

print(torch_tensor.dtype)
# torch.float32
print(np_array.dtype)
# float64
print(np_to_torch.dtype)
# torch.float64
```

This might lead to unwanted behaviour. I understand that the whole library is probably built around casting from numpy to other frameworks, so this might be difficult to solve. Perhaps `set_format` should include a `dtypes` option where for each input column the user can specify the wanted precision.

The alternative is that the user needs to cast manually after loading data from the dataset but that does not seem user-friendly, makes the dataset less portable, and might use more space in memory as well as on disk than is actually needed.

Comments: Indeed we convert tensors to list to be able to write in arrow format. Because of this conversion we lose the dtype information. We should add the dtype detection when we do type inference. However it would require a bit of refactoring since currently the conversion happens before the type inference..

And then for your information, when reading from arrow format we have to cast from arrow to numpy (which is fast since pyarrow has a numpy integration), and then to torch.

However there's one thing that can help you: we make sure that the dtypes correspond to what is defined in `features`.
Therefore what you can do is provide `features` in `.map(preprocess, feature=...)` to specify the output types.

For example in your case:
```python
from datasets import Features, Value, Sequence

features = Features({
    "input_ids": Sequence(Value("int32")),
    "sembedding": Sequence(Value("float32"))
})
preprocessed_dataset = dataset.map(preprocess, features=features)

preprocessed_dataset.set_format("torch", columns=["input_ids", "sembedding"])
print(preprocessed_dataset[0]["sembedding"].dtype)
# "torch.float32"
```

Let me know if it helps

If the arrow format is basically lists, why is the intermediate step to numpy necessary? I am a bit confused about that part.

Thanks for your suggestion. as I have currently implemented this, I cast to torch.Tensor in my collate_fn to save disk space (so I do not have to save padded tensors to max_len but can pad up to max batch len in collate_fn) at the cost of a bit slower processing. So for me this is not relevant anymore, but I am sure it is for others!

I'm glad you managed to figure something out :)

Casting from arrow to numpy can be 100x faster than casting from arrow to list.
This is because arrow has an integration with numpy that allows it to instantiate numpy arrays with zero-copy from arrow.
On the other hand to create python lists it is slow since it has to recreate the list object by iterating through each element in python.

Ah that is interesting. I have no direct experience with arrow so I didn't know. 

I encountered a simliar issue: `datasets` converted my float numpy array to `torch.float64` tensors, while many pytorch operations require `torch.float32` inputs and it's very troublesome. 

I tried @lhoestq 's solution, but since it's mixed with the preprocess function, it's not very intuitive. 

I just want to share another possible simpler solution: directly cast the dtype of the processed dataset.

Now I want to change the type of `labels` in `train_dataset` from float64 to float32, I can do this.

```
from datasets import Value, Sequence, Features
feats = train_dataset.features.copy()
feats['labels'].feature = Value(dtype='float32')
feats = Features(feats)
train_dataset.cast_(feats)
```


Reopening since @bhavitvyamalik started looking into it !

Also I'm posting here a function that could be helpful to support preserving the dtype of tensors.

It's used to build a pyarrow array out of a numpy array and:
- it doesn't convert the numpy array to a python list
- it keeps the precision of the numpy array for the pyarrow array
- it works with multidimensional arrays (while `pa.array` can only take a 1D array as input)
- it builds the pyarrow ListArray from offsets created on-the-fly and values that come from the flattened numpy array

```python
from functools import reduce
from operator import mul

import numpy as np
import pyarrow as pa

def pa_ndarray(a):
    """Build a PyArrow ListArray from a multidimensional NumPy array"""
    values = pa.array(a.flatten()) 
    for i in range(a.ndim - 1): 
        n_offsets = reduce(mul, a.shape[:a.ndim - i - 1], 1) 
        step_offsets = a.shape[a.ndim - i - 1] 
        offsets = pa.array(np.arange(n_offsets + 1) * step_offsets, type=pa.int32()) 
        values = pa.ListArray.from_arrays(offsets, values) 
    return values 

narr = np.arange(42).reshape(7, 2, 3).astype(np.uint8)
parr = pa_ndarray(narr)
assert isinstance(parr, pa.Array)
assert parr.type == pa.list_(pa.list_(pa.uint8()))
assert narr.tolist() == parr.to_pylist()
```

The only costly operation is the offsets computations. Since it doesn't iterate on the numpy array values this function is pretty fast.

@lhoestq Have you thought about this further?

We have a use case where we're attempting to load data containing numpy arrays using the `datasets` library.

When using one of the "standard" methods (`[Value(...)]` or `Sequence()`) we see ~200 samples processed per second during the call to `_prepare_split`. This slowdown is caused by the vast number of calls to `encode_nested_example` (each sequence is converted to a list, and each element in the sequence...). 

Using the `Feature` `ArrayND` improves this somewhat to ~500/s as it now uses numpy's `tolist()` rather than iterating over each value in the array and converting them individually.

However, it's still pretty slow and in theory it should be possible to avoid the `numpy -> python -> arrow` dance altogether. To demonstrate this, if you keep the `Feature` set to an `ArrayND` but instead return a `pa_ndarray(...)` in `_generate_examples` it skips the conversion (`return obj, False`) and hits ~11_000/s. Two orders of magnitude speed up! The problem is this then fails later on when the `ArrowWriter` tries to write the examples to disk :-( 

It would be nice to have first-class support for user-defined PyArrow objects. Is this a possibility? We have _large_ datasets where even an order of magnitude difference is important so settling on the middle ~500/s is less than ideal! 

Is there a workaround for this or another method that should be used instead that gets  near-to or equal performance to returning PyArrow arrays?

Note that manually generating the table using `pyarrow` achieves ~30_000/s

Hi !

It would be awesome to achieve this speed for numpy arrays !
For now we have to use `encode_nested_example` to convert numpy arrays to python lists since pyarrow doesn't support multidimensional numpy arrays (only 1D).

Maybe let's start a new PR from your PR @bhavitvyamalik (idk why we didn't answer your PR at that time, sorry about that).
Basically the idea is to allow `TypedSequence` to support numpy arrays as you did, and remove the numpy->python casting in `_cast_to_python_objects`.

This is really important since we are starting to have a focus on other modalities than text as well (audio, images).

Though until then @samgd, there is another feature that may interest you and that may give you the speed you want:

In a dataset script you can subclass either a GeneratorBasedBuilder (with the `_generate_examples ` method) or an ArrowBasedBuilder if you want. the ArrowBasedBuilder allows to yield arrow data by implementing the `_generate_tables` method (it's the same as `_generate_examples` except you must yield arrow tables). Since the data are already in arrow format, it doesn't call `encode_nested_example`. Let me know if that helps.