Forcing empty list/empty tuple behavior. (#514)

* Forcing empty list/empty tuple behavior. * Tailscale. * No tailscale anymore ? * Do not push the package it's not necessary. * Remove `Debug`.
huggingface · Aug 5, 2024 · aa4ad82 · aa4ad82
1 parent 6f791c5
commit aa4ad82
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 15 deletions.
diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml
@@ -95,9 +95,6 @@ jobs:
         uses: docker/setup-qemu-action@v2
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2
-      - uses: tailscale/github-action@v1
-        with:
-          authkey: ${{ secrets.TAILSCALE_AUTHKEY }}
       - name: Set short sha
         id: vars
         run: echo "GITHUB_SHA_SHORT=$(git rev-parse --short HEAD)" >> $GITHUB_ENV
@@ -107,7 +104,7 @@ jobs:
         with:
           # list of Docker images to use as base name for tags
           images: |
-            registry.internal.huggingface.tech/safetensors/s390x
+            ghcr.io/safetensors/s390x
           # generate Docker tags based on the following events/attributes
           tags: |
             type=schedule
@@ -118,18 +115,17 @@ jobs:
             type=semver,pattern={{major}}
             type=sha
       - name: Login to Registry
-        uses: docker/login-action@v2
+        uses: docker/login-action@v3
         with:
-          registry: registry.internal.huggingface.tech
-          username: ${{ secrets.REGISTRY_USERNAME }}
-          password: ${{ secrets.REGISTRY_PASSWORD }}
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
       - name: Test big endian
         uses: docker/build-push-action@v4
         with:
-          push: true
           platforms: linux/s390x
           file: Dockerfile.s390x.test
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
-          cache-from: type=registry,ref=registry.internal.huggingface.tech/safetensors/s390x:cache,mode=min
-          cache-to: type=registry,ref=registry.internal.huggingface.tech/safetensors/s390x:cache,mode=min
+          cache-from: type=registry,ref=ghcr.io/safetensors/s390x:cache,mode=max
+          cache-to: type=registry,ref=ghcr.io/safetensors/s390x:cache,mode=max
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
@@ -5,8 +5,7 @@ use pyo3::exceptions::{PyException, PyFileNotFoundError};
 use pyo3::prelude::*;
 use pyo3::sync::GILOnceCell;
 use pyo3::types::IntoPyDict;
-use pyo3::types::PySlice;
-use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList};
+use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList, PySlice};
 use pyo3::Bound as PyBound;
 use pyo3::{intern, PyErr};
 use safetensors::slice::TensorIndexer;
@@ -842,10 +841,27 @@ impl PySafeSlice {
     pub fn __getitem__(&self, slices: &PyBound<'_, PyAny>) -> PyResult<PyObject> {
         match &self.storage.as_ref() {
             Storage::Mmap(mmap) => {
-                let slices: Slice = slices.extract()?;
+                let pyslices = slices;
+                let slices: Slice = pyslices.extract()?;
+                let is_list = pyslices.is_instance_of::<PyList>();
                 let slices: Vec<SliceIndex> = match slices {
                     Slice::Slice(slice) => vec![slice],
-                    Slice::Slices(slices) => slices,
+                    Slice::Slices(slices) => {
+                        if slices.is_empty() && is_list {
+                            vec![SliceIndex::Slice(PySlice::new_bound(
+                                pyslices.py(),
+                                0,
+                                0,
+                                0,
+                            ))]
+                        } else if is_list {
+                            return Err(SafetensorError::new_err(
+                                "Non empty lists are not implemented",
+                            ));
+                        } else {
+                            slices
+                        }
+                    }
                 };
                 let data = &mmap[self.info.data_offsets.0 + self.offset
                     ..self.info.data_offsets.1 + self.offset];

diff --git a/bindings/python/tests/test_simple.py b/bindings/python/tests/test_simple.py
@@ -246,6 +246,10 @@ def test_torch_slice(self):
             self.assertEqual(list(tensor.shape), [10, 5])
             torch.testing.assert_close(tensor, A)
 
+            tensor = slice_[tuple()]
+            self.assertEqual(list(tensor.shape), [10, 5])
+            torch.testing.assert_close(tensor, A)
+
             tensor = slice_[:2]
             self.assertEqual(list(tensor.shape), [2, 5])
             torch.testing.assert_close(tensor, A[:2])
@@ -270,6 +274,10 @@ def test_torch_slice(self):
             self.assertEqual(list(tensor.shape), [8])
             torch.testing.assert_close(tensor, A[2:, -1])
 
+            tensor = slice_[list()]
+            self.assertEqual(list(tensor.shape), [0, 5])
+            torch.testing.assert_close(tensor, A[list()])
+
     def test_numpy_slice(self):
         A = np.random.rand(10, 5)
         tensors = {
@@ -284,6 +292,10 @@ def test_numpy_slice(self):
             self.assertEqual(list(tensor.shape), [10, 5])
             self.assertTrue(np.allclose(tensor, A))
 
+            tensor = slice_[tuple()]
+            self.assertEqual(list(tensor.shape), [10, 5])
+            self.assertTrue(np.allclose(tensor, A))
+
             tensor = slice_[:2]
             self.assertEqual(list(tensor.shape), [2, 5])
             self.assertTrue(np.allclose(tensor, A[:2]))
@@ -312,10 +324,18 @@ def test_numpy_slice(self):
             self.assertEqual(list(tensor.shape), [8])
             self.assertTrue(np.allclose(tensor, A[2:, -5]))
 
+            tensor = slice_[list()]
+            self.assertEqual(list(tensor.shape), [0, 5])
+            self.assertTrue(np.allclose(tensor, A[list()]))
+
             with self.assertRaises(SafetensorError) as cm:
                 tensor = slice_[2:, -6]
             self.assertEqual(str(cm.exception), "Invalid index -6 for dimension 1 of size 5")
 
+            with self.assertRaises(SafetensorError) as cm:
+                tensor = slice_[[0, 1]]
+            self.assertEqual(str(cm.exception), "Non empty lists are not implemented")
+
             with self.assertRaises(SafetensorError) as cm:
                 tensor = slice_[2:, 20]
             self.assertEqual(