From 922bbb784352523cf39fc07a0d947f5e33ee0968 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=BE=B7=E9=92=8A?= <lidezhao1@huawei.com>
Date: Tue, 29 Apr 2025 10:35:02 +0800
Subject: [PATCH] add drop_out_v3

---
 tf_adapter/kernels/aicore/npu_aicore_ops.cc   | 10 ++++++++++
 tf_adapter/ops/aicore/npu_aicore_ops.cc       | 14 ++++++++++++++
 .../python/npu_bridge/npu_cpu/npu_cpu_ops.py  | 19 +++++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/tf_adapter/kernels/aicore/npu_aicore_ops.cc b/tf_adapter/kernels/aicore/npu_aicore_ops.cc
index 83a027b6e..b66740143 100644
--- a/tf_adapter/kernels/aicore/npu_aicore_ops.cc
+++ b/tf_adapter/kernels/aicore/npu_aicore_ops.cc
@@ -91,6 +91,16 @@ class FastGeluGradOp : public tensorflow::OpKernel {
   }
 };
 
+class StatelessDropoutOp : public tensorflow::OpKernel {
+public:
+  explicit StatelessDropoutOp(tensorflow::OpKernelConstruction *context) : OpKernel(context) {}
+  ~StatelessDropoutOp() override {}
+  void Compute(tensorflow::OpKernelContext *context) override {}
+};
+
+REGISTER_KERNEL_BUILDER(Name("StatelessDropout")
+.Device(tensorflow::DEVICE_CPU), StatelessDropoutOp);
+
 REGISTER_KERNEL_BUILDER(
     Name("FastGeluGrad")
 .
diff --git a/tf_adapter/ops/aicore/npu_aicore_ops.cc b/tf_adapter/ops/aicore/npu_aicore_ops.cc
index fbe6035bd..dfaac8e8f 100644
--- a/tf_adapter/ops/aicore/npu_aicore_ops.cc
+++ b/tf_adapter/ops/aicore/npu_aicore_ops.cc
@@ -97,6 +97,20 @@ REGISTER_OP("DynamicGruV2")
       return Status::OK();
     });
 
+REGISTER_OP("StatelessDropout")
+    .Input("x: T")
+    .Input("noise_shape: int64")
+    .Input("p: T")
+    .Input("seed: int64")
+    .Input("offset: int64")
+    .Output("y: T")
+    .Attr("T: {float16, float32, bfloat16}")
+    .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext *c) {
+      c->set_output(0, c->input(0));
+      return Status::OK();
+    });
+
 REGISTER_OP("DynamicGruV2Grad")
     .Input("x: T")
     .Input("weight_input: T")
diff --git a/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py b/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py
index 417119584..ba6ed2b5d 100644
--- a/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py
+++ b/tf_adapter/python/npu_bridge/npu_cpu/npu_cpu_ops.py
@@ -87,6 +87,25 @@ def dense_image_warp(image, flow, name=None):
     return result
 
 
+## 提供host侧StatelessDropout功能
+#  @param x float32,float16,bfloat16 类型
+#  @param noise_shape int64 类型
+#  @param p float32,float16,bfloat16 类型
+#  @param seed int64 类型
+#  @param offset int64 类型
+#  @return values float32,float16,bfloat16 类型
+def stateless_dropout(x, noise_shape, p, seed, offset):
+    """ host stateless_dropout. """
+    result = gen_npu_cpu_ops.StatelessDropout(
+        x=x,
+        noise_shape=noise_shape,
+        p=p,
+        seed=seed,
+        offset=offset   
+    )
+    return result
+
+
 ## DenseImageWarp的梯度函数
 @ops.RegisterGradient("DenseImageWarp")
 def dense_image_warp_grad(op, grad):
-- 
Gitee