chore: import upstream snapshot with attribution

tensorflow/python/dlpack/BUILD

@@ -0,0 +1,36 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:private"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "dlpack",
+    srcs = ["dlpack.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python:pywrap_tfe",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+cuda_py_strict_test(
+    name = "dlpack_test",
+    srcs = ["dlpack_test.py"],
+    deps = [
+        ":dlpack",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/dlpack/dlpack.py

@@ -0,0 +1,63 @@
+# 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.
+# ==============================================================================
+"""DLPack modules for Tensorflow."""
+
+from tensorflow.python import pywrap_tfe
+from tensorflow.python.eager import context
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("experimental.dlpack.to_dlpack", v1=[])
+def to_dlpack(tf_tensor):
+  """Returns the dlpack capsule representing the tensor.
+
+  This operation ensures the underlying data memory is ready when returns.
+
+    ```python
+    a = tf.tensor([1, 10])
+    dlcapsule = tf.experimental.dlpack.to_dlpack(a)
+    # dlcapsule represents the dlpack data structure
+    ```
+
+  Args:
+    tf_tensor: Tensorflow eager tensor, to be converted to dlpack capsule.
+
+  Returns:
+    A PyCapsule named as dltensor, which shares the underlying memory to other
+     framework. This PyCapsule can be consumed only once.
+  """
+  return pywrap_tfe.TFE_ToDlpackCapsule(tf_tensor)
+
+
+@tf_export("experimental.dlpack.from_dlpack", v1=[])
+def from_dlpack(dlcapsule):
+  """Returns the Tensorflow eager tensor.
+
+  The returned tensor uses the memory shared by dlpack capsules from other
+  framework.
+
+    ```python
+    a = tf.experimental.dlpack.from_dlpack(dlcapsule)
+    # `a` uses the memory shared by dlpack
+    ```
+
+  Args:
+    dlcapsule: A PyCapsule named as dltensor
+
+  Returns:
+    A Tensorflow eager tensor
+  """
+  context.context().ensure_initialized()
+  return pywrap_tfe.TFE_FromDlpackCapsule(dlcapsule, context.context()._handle)  # pylint: disable=protected-access

tensorflow/python/dlpack/dlpack_test.py

@@ -0,0 +1,129 @@
+# 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 DLPack functions."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.dlpack import dlpack
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+int_dtypes = [
+    np.int8, np.int16, np.int32, np.int64, np.uint8, np.uint16, np.uint32,
+    np.uint64
+]
+float_dtypes = [np.float16, np.float32, np.float64]
+complex_dtypes = [np.complex64, np.complex128]
+dlpack_dtypes = (
+    int_dtypes + float_dtypes + [dtypes.bfloat16] + complex_dtypes + [np.bool_]
+)
+
+testcase_shapes = [(), (1,), (2, 3), (2, 0), (0, 7), (4, 1, 2)]
+
+
+def FormatShapeAndDtype(shape, dtype):
+  return "_{}[{}]".format(str(dtype), ",".join(map(str, shape)))
+
+
+def GetNamedTestParameters():
+  result = []
+  for dtype in dlpack_dtypes:
+    for shape in testcase_shapes:
+      result.append({
+          "testcase_name": FormatShapeAndDtype(shape, dtype),
+          "dtype": dtype,
+          "shape": shape
+      })
+  return result
+
+
+class DLPackTest(parameterized.TestCase, test.TestCase):
+
+  @parameterized.named_parameters(GetNamedTestParameters())
+  def testRoundTrip(self, dtype, shape):
+    np.random.seed(42)
+    if dtype == np.bool_:
+      np_array = np.random.randint(0, 1, shape, np.bool_)
+    else:
+      np_array = np.random.randint(0, 10, shape)
+    # copy to gpu if available
+    tf_tensor = array_ops.identity(constant_op.constant(np_array, dtype=dtype))
+    tf_tensor_device = tf_tensor.device
+    tf_tensor_dtype = tf_tensor.dtype
+    dlcapsule = dlpack.to_dlpack(tf_tensor)
+    del tf_tensor  # should still work
+    tf_tensor2 = dlpack.from_dlpack(dlcapsule)
+    self.assertAllClose(np_array, tf_tensor2)
+    if tf_tensor_dtype == dtypes.int32:
+      # int32 tensor is always on cpu for now
+      self.assertEqual(tf_tensor2.device,
+                       "/job:localhost/replica:0/task:0/device:CPU:0")
+    else:
+      self.assertEqual(tf_tensor_device, tf_tensor2.device)
+
+  def testRoundTripWithoutToDlpack(self):
+    np_array = np.random.randint(0, 10, [42])
+    self.assertAllEqual(
+        np.from_dlpack(constant_op.constant(np_array).cpu()), np_array
+    )
+
+  def testTensorsCanBeConsumedOnceOnly(self):
+    np.random.seed(42)
+    np_array = np.random.randint(0, 10, (2, 3, 4))
+    tf_tensor = constant_op.constant(np_array, dtype=np.float32)
+    dlcapsule = dlpack.to_dlpack(tf_tensor)
+    del tf_tensor  # should still work
+    _ = dlpack.from_dlpack(dlcapsule)
+
+    def ConsumeDLPackTensor():
+      dlpack.from_dlpack(dlcapsule)  # Should can be consumed only once
+
+    self.assertRaisesRegex(Exception,
+                           ".*a DLPack tensor may be consumed at most once.*",
+                           ConsumeDLPackTensor)
+
+  def testDLPackFromWithoutContextInitialization(self):
+    tf_tensor = constant_op.constant(1)
+    dlcapsule = dlpack.to_dlpack(tf_tensor)
+    # Resetting the context doesn't cause an error.
+    context._reset_context()
+    _ = dlpack.from_dlpack(dlcapsule)
+
+  def testUnsupportedTypeToDLPack(self):
+
+    def UnsupportedQint16():
+      tf_tensor = constant_op.constant([[1, 4], [5, 2]], dtype=dtypes.qint16)
+      _ = dlpack.to_dlpack(tf_tensor)
+
+    self.assertRaisesRegex(Exception, ".* is not supported by dlpack",
+                           UnsupportedQint16)
+
+  def testMustPassTensorArgumentToDLPack(self):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "The argument to `to_dlpack` must be a TF tensor, not Python object"):
+      dlpack.to_dlpack([1])
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()