Update docstrings & type-hints (#139)

huggingface · Dec 27, 2022 · ba8de5f · ba8de5f
1 parent a104c5e
commit ba8de5f
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/bindings/python/py_src/safetensors/flax.py b/bindings/python/py_src/safetensors/flax.py
@@ -25,7 +25,7 @@ def save(tensors: Dict[str, jnp.DeviceArray], metadata: Optional[Dict[str, str]]
 
     ```python
     from safetensors.flax import save
-    import flax
+    from jax import numpy as jnp
 
     tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))}
     byte_data = save(tensors)
@@ -39,7 +39,7 @@ def save_file(
     tensors: Dict[str, jnp.DeviceArray],
     filename: str,
     metadata: Optional[Dict[str, str]] = None,
-):
+) -> None:
     """
     Saves a dictionnary of tensors into raw bytes in safetensors format.
 
@@ -60,7 +60,7 @@ def save_file(
 
     ```python
     from safetensors.flax import save_file
-    import flax
+    from jax import numpy as jnp
 
     tensors = {"embedding": jnp.zeros((512, 1024)), "attention": jnp.zeros((256, 256))}
     save(tensors, "model.safetensors")

diff --git a/bindings/python/py_src/safetensors/numpy.py b/bindings/python/py_src/safetensors/numpy.py
@@ -41,7 +41,7 @@ def save(tensor_dict: Dict[str, np.ndarray], metadata: Optional[Dict[str, str]]
     return result
 
 
-def save_file(tensor_dict: Dict[str, np.ndarray], filename: str, metadata: Optional[Dict[str, str]] = None):
+def save_file(tensor_dict: Dict[str, np.ndarray], filename: str, metadata: Optional[Dict[str, str]] = None) -> None:
     """
     Saves a dictionnary of tensors into raw bytes in safetensors format.
 

diff --git a/bindings/python/py_src/safetensors/tensorflow.py b/bindings/python/py_src/safetensors/tensorflow.py
@@ -39,7 +39,7 @@ def save_file(
     tensors: Dict[str, tf.Tensor],
     filename: str,
     metadata: Optional[Dict[str, str]] = None,
-):
+) -> None:
     """
     Saves a dictionnary of tensors into raw bytes in safetensors format.