From ef5175344fa58d7387ccb8bf80dee6f44b833813 Mon Sep 17 00:00:00 2001
From: JiangDongHua <759421566@qq.com>
Date: Wed, 21 Aug 2024 22:18:42 +0800
Subject: [PATCH] add Tensor method: new and new_ones

---
 python/oneflow/__init__.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/python/oneflow/__init__.py b/python/oneflow/__init__.py
index 8706117f6a8..c51c1dc3444 100644
--- a/python/oneflow/__init__.py
+++ b/python/oneflow/__init__.py
@@ -508,3 +508,13 @@ def atexit_hook(hook):
 if oneflow._oneflow_internal.flags.with_mlir():
     oneflow_internal_path = oneflow._oneflow_internal.__file__
     oneflow._oneflow_internal.ir.load_jit_shared_lib(oneflow_internal_path)
+
+def flow_ones(self, *args, **kwargs):
+    return ones(*args, **kwargs, device=self.device, dtype=self.dtype)
+
+Tensor.new_ones = flow_ones
+
+def flow_zeros(self, *args, **kwargs):
+    return zeros(*args, **kwargs, device=self.device, dtype=self.dtype)
+
+Tensor.new = flow_zeros