This PR fixes some warnings, failing tests, but primarily it's meant to change the behavior of train and predict to accept Nx containers. After discussing with @josevalim we decided that it was best to use Nx.to_tensor on all inputs so people could implement the Nx container protocol and this library just works with anything compatible with Nx (such as Explorer). After messing with the test I realized concatenate is probably better because all containers will almost always be collections of tensors ---given train and predict require collections (e.g. batches) of inputs. The single concat has no affect on input tensors.
This PR fixes some warnings, failing tests, but primarily it's meant to change the behavior of train and predict to accept Nx containers. After discussing with @josevalim we decided that it was best to use Nx.to_tensor on all inputs so people could implement the Nx container protocol and this library just works with anything compatible with Nx (such as Explorer). After messing with the test I realized concatenate is probably better because all containers will almost always be collections of tensors ---given train and predict require collections (e.g. batches) of inputs. The single concat has no affect on input tensors.