zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

update: format code

Browse Source
This commit is contained in:
John Mai
2025-06-15 17:35:33 +08:00
parent 989e8bab66
commit b3c1aaafd2
2 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
benchmarks/python/comparative/bench_mlx.py
View File

@@ -223,12 +223,14 @@ def relu6(x):
y = nn.relu6(y) y = nn.relu6(y)
mx.eval(y) mx.eval(y)
def relu_squared(x): def relu_squared(x):
y = x y = x
for i in range(100): for i in range(100):
y = nn.relu_squared(y) y = nn.relu_squared(y)
mx.eval(y) mx.eval(y)
def softplus(x): def softplus(x):
y = x y = x
for i in range(100): for i in range(100):

2
benchmarks/python/comparative/bench_torch.py
View File

@@ -156,6 +156,7 @@ def relu6(x):
y = torch.nn.functional.relu6(y) y = torch.nn.functional.relu6(y)
sync_if_needed(x) sync_if_needed(x)
@torch.no_grad() @torch.no_grad()
def relu_squared(x): def relu_squared(x):
y = x y = x
@@ -164,6 +165,7 @@ def relu_squared(x):
y = torch.square(y) y = torch.square(y)
sync_if_needed(x) sync_if_needed(x)
@torch.no_grad() @torch.no_grad()
def softplus(x): def softplus(x):
y = x y = x