NVIDIA/cutlass
1
0
Fork 0
mirror of https://github.com/NVIDIA/cutlass.git synced 2026-04-19 22:38:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2919 from pbelevich/patch-1

Browse Source 
Refactor binary_op functions to remove unused result parameter
This commit is contained in:
drazi
2026-02-11 11:48:58 +08:00
committed by GitHub
parent 6b3e607b85 b6d7703e02
commit 71aa7a0abc
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/python/CuTeDSL/notebooks/tensorssa.ipynb
View File

@@ -167,7 +167,7 @@
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_1(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
"def binary_op_1(a: cute.Tensor, b: cute.Tensor):\n",
" a_vec = a.load()\n",
" b_vec = b.load()\n",
"\n",
@@ -184,7 +184,7 @@
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n",
" floor_div_res = a_vec // b_vec\n",
" cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n",
" cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n",
" mod_res = a_vec % b_vec\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
@@ -194,8 +194,7 @@
"a.fill(1.0)\n",
"b = np.empty((3,), dtype=np.float32)\n",
"b.fill(2.0)\n",
"res = np.empty((3,), dtype=np.float32)\n",
"binary_op_1(from_dlpack(res), from_dlpack(a), from_dlpack(b))"
"binary_op_1(from_dlpack(a), from_dlpack(b))"
]
},
{
@@ -205,7 +204,7 @@
"outputs": [],
"source": [
"@cute.jit\n",
"def binary_op_2(res: cute.Tensor, a: cute.Tensor, c: cutlass.Constexpr):\n",
"def binary_op_2(a: cute.Tensor, c: cutlass.Constexpr):\n",
" a_vec = a.load()\n",
"\n",
" add_res = a_vec + c\n",
@@ -230,8 +229,7 @@
"a = np.empty((3,), dtype=np.float32)\n",
"a.fill(1.0)\n",
"c = 2.0\n",
"res = np.empty((3,), dtype=np.float32)\n",
"binary_op_2(from_dlpack(res), from_dlpack(a), c)"
"binary_op_2(from_dlpack(a), c)"
]
},
{