From 51cdbb093701d355e5f7e036e4d9a70ef2ad17ca Mon Sep 17 00:00:00 2001
From: Hoel Bagard <hoel.bagard@gmail.com>
Date: Fri, 17 Oct 2025 14:05:20 +0900
Subject: [PATCH] tf: add tf.split

---
 stubs/tensorflow/tensorflow/__init__.pyi | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/stubs/tensorflow/tensorflow/__init__.pyi b/stubs/tensorflow/tensorflow/__init__.pyi
index 3a356392a6aa..3124c2c6e556 100644
--- a/stubs/tensorflow/tensorflow/__init__.pyi
+++ b/stubs/tensorflow/tensorflow/__init__.pyi
@@ -385,6 +385,13 @@ def squeeze(
 ) -> Tensor: ...
 @overload
 def squeeze(input: RaggedTensor, axis: int | tuple[int, ...] | list[int], name: str | None = None) -> RaggedTensor: ...
+def split(
+    value: TensorCompatible,
+    num_or_size_splits: int | TensorCompatible,
+    axis: int | Tensor = 0,
+    num: int | None = None,
+    name: str | None = "split",
+) -> list[Tensor]: ...
 def tensor_scatter_nd_update(
     tensor: TensorCompatible, indices: TensorCompatible, updates: TensorCompatible, name: str | None = None
 ) -> Tensor: ...