add Tensor method: new and new_ones

Oneflow-Inc · Aug 21, 2024 · ef51753 · ef51753
1 parent f0d13a6
commit ef51753
Showing 1 changed file with 10 additions and 0 deletions.
diff --git 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