add tensorflow.gather_nd

python · Jan 5, 2025 · 3f6d9b1 · 3f6d9b1
1 parent 4e678f2
commit 3f6d9b1
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -20,7 +20,15 @@ from tensorflow import (
     math as math,
     types as types,
 )
-from tensorflow._aliases import AnyArray, DTypeLike, ScalarTensorCompatible, ShapeLike, Slice, TensorCompatible
+from tensorflow._aliases import (
+    AnyArray,
+    DTypeLike,
+    ScalarTensorCompatible,
+    ShapeLike,
+    Slice,
+    TensorCompatible,
+    UIntTensorCompatible,
+)
 from tensorflow.autodiff import GradientTape as GradientTape
 from tensorflow.core.protobuf import struct_pb2
 from tensorflow.dtypes import *
@@ -415,4 +423,7 @@ def shape(input: TensorCompatible, out_type: DTypeLike | None = None, name: str
 def where(
     condition: TensorCompatible, x: TensorCompatible | None = None, y: TensorCompatible | None = None, name: str | None = None
 ) -> Tensor: ...
+def gather_nd(
+    params: TensorCompatible, indices: UIntTensorCompatible, batch_dims: UIntTensorCompatible = 0, name: str | None = None
+) -> Tensor: ...
 def __getattr__(name: str) -> Incomplete: ...