From 18995cad1cf3e08b0675859544aefb7f64904478 Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 18 Feb 2024 20:35:44 +0900 Subject: [PATCH 1/4] add pre-existing stubs. --- stubs/tensorflow/tensorflow/bitwise.pyi | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 stubs/tensorflow/tensorflow/bitwise.pyi diff --git a/stubs/tensorflow/tensorflow/bitwise.pyi b/stubs/tensorflow/tensorflow/bitwise.pyi new file mode 100644 index 000000000000..8b886826b55b --- /dev/null +++ b/stubs/tensorflow/tensorflow/bitwise.pyi @@ -0,0 +1,13 @@ +from typing import overload + +import tensorflow as tf +from tensorflow._aliases import TensorCompatible + +@overload +def bitwise_or(x: TensorCompatible, y: TensorCompatible, name: str | None = None) -> tf.Tensor: ... +@overload +def bitwise_or(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... +@overload +def bitwise_and(x: TensorCompatible, y: TensorCompatible, name: str | None = None) -> tf.Tensor: ... +@overload +def bitwise_and(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... From 7051cdbd329b65ff71a620134c9d2f0da68e2fff Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sun, 18 Feb 2024 20:44:09 +0900 Subject: [PATCH 2/4] add missing functions. --- stubs/tensorflow/tensorflow/bitwise.pyi | 30 ++++++++++++++++++++----- 1 file changed, 25 insertions(+), 5 deletions(-) diff --git a/stubs/tensorflow/tensorflow/bitwise.pyi b/stubs/tensorflow/tensorflow/bitwise.pyi index 8b886826b55b..54daaf7f7db6 100644 --- a/stubs/tensorflow/tensorflow/bitwise.pyi +++ b/stubs/tensorflow/tensorflow/bitwise.pyi @@ -1,13 +1,33 @@ -from typing import overload +from typing import Any, overload +from typing_extensions import TypeAlias +import numpy as np import tensorflow as tf -from tensorflow._aliases import TensorCompatible +from tensorflow._aliases import FloatArray, IntArray +_BitwiseCompatible: TypeAlias = tf.Tensor | bool | int | FloatArray | IntArray | np.number[Any] + +@overload +def bitwise_and(x: _BitwiseCompatible, y: _BitwiseCompatible, name: str | None = None) -> tf.Tensor: ... +@overload +def bitwise_and(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... @overload -def bitwise_or(x: TensorCompatible, y: TensorCompatible, name: str | None = None) -> tf.Tensor: ... +def bitwise_or(x: _BitwiseCompatible, y: _BitwiseCompatible, name: str | None = None) -> tf.Tensor: ... @overload def bitwise_or(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... @overload -def bitwise_and(x: TensorCompatible, y: TensorCompatible, name: str | None = None) -> tf.Tensor: ... +def bitwise_xor(x: _BitwiseCompatible, y: _BitwiseCompatible, name: str | None = None) -> tf.Tensor: ... @overload -def bitwise_and(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... +def bitwise_xor(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... +@overload +def invert(x: _BitwiseCompatible, name: str | None = None) -> tf.Tensor: ... +@overload +def invert(x: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... +@overload +def left_shift(x: _BitwiseCompatible, y: _BitwiseCompatible, name: str | None = None) -> tf.Tensor: ... +@overload +def left_shift(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... +@overload +def right_shift(x: _BitwiseCompatible, y: _BitwiseCompatible, name: str | None = None) -> tf.Tensor: ... +@overload +def right_shift(x: tf.RaggedTensor, y: tf.RaggedTensor, name: str | None = None) -> tf.RaggedTensor: ... From 663e76a04bc36609f93ebf48867ca5173a8635df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ho=C3=ABl=20Bagard?= <34478245+hoel-bagard@users.noreply.github.com> Date: Sat, 2 Mar 2024 08:41:14 +0900 Subject: [PATCH 3/4] Update stubs/tensorflow/tensorflow/bitwise.pyi Co-authored-by: Jelle Zijlstra --- stubs/tensorflow/tensorflow/bitwise.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/stubs/tensorflow/tensorflow/bitwise.pyi b/stubs/tensorflow/tensorflow/bitwise.pyi index 54daaf7f7db6..631aa625b588 100644 --- a/stubs/tensorflow/tensorflow/bitwise.pyi +++ b/stubs/tensorflow/tensorflow/bitwise.pyi @@ -5,7 +5,7 @@ import numpy as np import tensorflow as tf from tensorflow._aliases import FloatArray, IntArray -_BitwiseCompatible: TypeAlias = tf.Tensor | bool | int | FloatArray | IntArray | np.number[Any] +_BitwiseCompatible: TypeAlias = tf.Tensor | int | FloatArray | IntArray | np.number[Any] @overload def bitwise_and(x: _BitwiseCompatible, y: _BitwiseCompatible, name: str | None = None) -> tf.Tensor: ... From b3d3798fc635c0fd065419f29ce892a13ae39dea Mon Sep 17 00:00:00 2001 From: Hoel Bagard Date: Sat, 2 Mar 2024 08:46:52 +0900 Subject: [PATCH 4/4] add comments about type alias not being fully accurate. --- stubs/tensorflow/tensorflow/bitwise.pyi | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/stubs/tensorflow/tensorflow/bitwise.pyi b/stubs/tensorflow/tensorflow/bitwise.pyi index 631aa625b588..2045d9987326 100644 --- a/stubs/tensorflow/tensorflow/bitwise.pyi +++ b/stubs/tensorflow/tensorflow/bitwise.pyi @@ -5,6 +5,10 @@ import numpy as np import tensorflow as tf from tensorflow._aliases import FloatArray, IntArray +# The alias below is not fully accurate, since TensorFlow casts the inputs, they have some additional +# requirements. For example y needs to be castable into x's dtype. Moreover, x and y cannot both be booleans. +# Properly typing the bitwise functions would be overly complicated and unlikely to provide much benefits +# since most people use Tensors, it was therefore not done. _BitwiseCompatible: TypeAlias = tf.Tensor | int | FloatArray | IntArray | np.number[Any] @overload