diff --git a/tf_adapter/kernels/aicore/npu_aicore_ops.cc b/tf_adapter/kernels/aicore/npu_aicore_ops.cc index 83a027b6ea2e367f2d36cf86c431216bc6bcc8e5..b6674014350bad0613a34b50815b950ac5e05829 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 fbe6035bdb24c97067cf9517a9dd736a31070c3d..dfaac8e8f66b06e3564e9ef236e987025fb21b35 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 417119584f09363575c133283ad275aab85380dd..ba6ed2b5dcd4a6305f968ec4081686a7a6838670 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):