chore: import upstream snapshot with attribution

tensorflow/lite/tools/optimize/sparsity/BUILD

@@ -0,0 +1,50 @@
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
+load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+pybind_extension(
+    name = "format_converter_wrapper_pybind11",
+    srcs = ["format_converter_wrapper_pybind11.cc"],
+    common_lib_packages = [
+        "litert/python",
+        "tensorflow/lite/python",
+    ],
+    copts = [
+        "-fexceptions",
+        "-fno-strict-aliasing",
+    ],
+    data = [
+        "format_converter_wrapper_pybind11.pyi",
+    ],
+    enable_stub_generation = True,
+    features = ["-use_header_modules"],
+    wrap_py_init = True,
+    deps = [
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels/internal/utils:sparsity_format_converter",
+        "@pybind11",
+        "@xla//third_party/python_runtime:headers",  # buildcleaner: keep
+    ],
+)
+
+py_test(
+    name = "format_converter_wrapper_pybind11_test",
+    srcs = ["format_converter_wrapper_pybind11_test.py"],
+    strict_deps = True,
+    deps = [
+        ":format_converter_wrapper_pybind11",
+        "@absl_py//absl/testing:absltest",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+    ] + if_pywrap(
+        if_true = ["//tensorflow/lite/python:pywrap_tflite"],
+    ),
+)

tensorflow/lite/tools/optimize/sparsity/format_converter_wrapper_pybind11.cc

@@ -0,0 +1,64 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <vector>
+
+#include "pybind11/attr.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/utils/sparsity_format_converter.h"
+
+namespace py = pybind11;
+
+using FormatConverterFp32 = tflite::internal::sparsity::FormatConverter<float>;
+
+PYBIND11_MODULE(format_converter_wrapper_pybind11, m) {
+  m.doc() = "Python wrapper for the tflite sparse tensor converter.";
+
+  py::enum_<TfLiteStatus>(m, "TfLiteStatus")
+      .value("TF_LITE_OK", TfLiteStatus::kTfLiteOk)
+      .value("TF_LITE_ERROR", TfLiteStatus::kTfLiteError)
+      .export_values();
+
+  py::enum_<TfLiteDimensionType>(m, "TfLiteDimensionType")
+      .value("TF_LITE_DIM_DENSE", TfLiteDimensionType::kTfLiteDimDense)
+      .value("TF_LITE_DIM_SPARSE_CSR", TfLiteDimensionType::kTfLiteDimSparseCSR)
+      .export_values();
+
+  py::class_<TfLiteSparsity> sparsity_class(m, "TfLiteSparsity",
+                                            py::module_local());
+
+  py::class_<FormatConverterFp32>(m, "FormatConverterFp32")
+      .def(py::init</*shape=*/const std::vector<int>&,
+                    /*traversal_order=*/const std::vector<int>&,
+                    /*format=*/const std::vector<TfLiteDimensionType>&,
+                    /*block_size=*/const std::vector<int>&,
+                    /*block_map=*/const std::vector<int>&>())
+      .def(py::init</*shape=*/const std::vector<int>&,
+                    /*sparsity=*/const TfLiteSparsity&>())
+      .def("GetData", &FormatConverterFp32::GetData)
+      .def("GetDimMetadata", &FormatConverterFp32::GetDimMetadata)
+      .def("DenseToSparse",
+           [](FormatConverterFp32& converter, py::buffer buf) {
+             py::buffer_info buffer_info = buf.request();
+             return converter.DenseToSparse(
+                 static_cast<float*>(buffer_info.ptr));
+           })
+      .def("SparseToDense", [](FormatConverterFp32& converter, py::buffer buf) {
+        py::buffer_info buffer_info = buf.request();
+        return converter.SparseToDense(static_cast<float*>(buffer_info.ptr));
+      });
+}

tensorflow/lite/tools/optimize/sparsity/format_converter_wrapper_pybind11.pyi

@@ -0,0 +1,66 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import ClassVar, overload
+
+TF_LITE_DIM_DENSE: TfLiteDimensionType
+TF_LITE_DIM_SPARSE_CSR: TfLiteDimensionType
+TF_LITE_ERROR: TfLiteStatus
+TF_LITE_OK: TfLiteStatus
+
+class FormatConverterFp32:
+    @overload
+    def __init__(self, arg0: list[int], arg1: list[int], arg2: list[TfLiteDimensionType], arg3: list[int], arg4: list[int]) -> None: ...
+    @overload
+    def __init__(self, arg0: list[int], arg1: TfLiteSparsity) -> None: ...
+    def DenseToSparse(self, arg0) -> TfLiteStatus: ...
+    def GetData(self) -> list[float]: ...
+    def GetDimMetadata(self) -> list[list[int]]: ...
+    def SparseToDense(self, arg0) -> TfLiteStatus: ...
+
+class TfLiteDimensionType:
+    __members__: ClassVar[dict] = ...  # read-only
+    TF_LITE_DIM_DENSE: ClassVar[TfLiteDimensionType] = ...
+    TF_LITE_DIM_SPARSE_CSR: ClassVar[TfLiteDimensionType] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...
+
+class TfLiteSparsity:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TfLiteStatus:
+    __members__: ClassVar[dict] = ...  # read-only
+    TF_LITE_ERROR: ClassVar[TfLiteStatus] = ...
+    TF_LITE_OK: ClassVar[TfLiteStatus] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...

tensorflow/lite/tools/optimize/sparsity/format_converter_wrapper_pybind11_test.py

@@ -0,0 +1,69 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Tests for the pybind11 bindings of format_converter."""
+
+from absl.testing import absltest
+import numpy as np
+
+from tensorflow.lite.tools.optimize.sparsity import format_converter_wrapper_pybind11 as format_converter
+
+
+class FormatConverterTest(absltest.TestCase):
+
+  def test_bcsr_fp32(self):
+    """Same as FormatConverterTest::BlockTestD0S1 but via pybind11."""
+    # pyformat: disable
+    dense_matrix = [1.0, 0.0, 2.0, 3.0,
+                    0.0, 4.0, 0.0, 0.0,
+                    0.0, 0.0, 5.0, 0.0,
+                    0.0, 0.0, 0.0, 6.0]
+    # pyformat: enable
+    dense_shape = [4, 4]
+    traversal_order = [0, 1, 2, 3]
+    dim_types = [
+        format_converter.TfLiteDimensionType.TF_LITE_DIM_DENSE,
+        format_converter.TfLiteDimensionType.TF_LITE_DIM_SPARSE_CSR
+    ]
+    block_size = [2, 2]
+    block_map = [0, 1]
+    converter = format_converter.FormatConverterFp32(dense_shape,
+                                                     traversal_order, dim_types,
+                                                     block_size, block_map)
+
+    converter.DenseToSparse(np.asarray(dense_matrix, dtype=np.float32).data)
+
+    dim_metadata = converter.GetDimMetadata()
+    self.assertEqual([2], dim_metadata[0])
+    self.assertEmpty(dim_metadata[1])  # rows are dense.
+
+    self.assertEqual([0, 2, 3], dim_metadata[2])  # array segments.
+    self.assertEqual([0, 1, 1], dim_metadata[3])  # array indices.
+
+    self.assertEqual([2], dim_metadata[4])
+    self.assertEmpty(dim_metadata[5])  # sub block rows are dense.
+
+    self.assertEqual([2], dim_metadata[6])
+    self.assertEmpty(dim_metadata[7])  # sub block columns are dense.
+
+    expected_data = [1.0, 0.0, 0.0, 4.0, 2.0, 3.0, 0.0, 0.0, 5.0, 0.0, 0.0, 6.0]
+    sparse_data = converter.GetData()
+    self.assertTrue(np.allclose(expected_data, sparse_data))
+
+    converter.SparseToDense(np.asarray(sparse_data, dtype=np.float32).data)
+    self.assertTrue(np.allclose(dense_matrix, converter.GetData()))
+
+
+if __name__ == '__main__':
+  absltest.main()