Framework-agnostic version of `numpy.size` that will work on torch/TensorFlow/Jax tensors as well as NumPy arrays.
(array)
| 692 | |
| 693 | |
| 694 | def tensor_size(array): |
| 695 | """ |
| 696 | Framework-agnostic version of `numpy.size` that will work on torch/TensorFlow/Jax tensors as well as NumPy arrays. |
| 697 | """ |
| 698 | if is_numpy_array(array): |
| 699 | return np.size(array) |
| 700 | elif is_torch_tensor(array): |
| 701 | return array.numel() |
| 702 | elif is_tf_tensor(array): |
| 703 | import tensorflow as tf |
| 704 | |
| 705 | return tf.size(array) |
| 706 | elif is_jax_tensor(array): |
| 707 | return array.size |
| 708 | else: |
| 709 | raise ValueError(f"Type not supported for tensor_size: {type(array)}.") |
| 710 | |
| 711 | |
| 712 | def add_model_info_to_auto_map(auto_map, repo_id): |
no test coverage detected