zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-13 11:41:14 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Match PT for nearest upsample

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-05-18 19:33:47 -07:00
parent 8576e6fe36
commit 3cdc4a5f70
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/mlx/nn/layers/upsample.py
View File

@ -12,7 +12,7 @@ from mlx.nn.layers.base import Module
def _scaled_indices(N, scale, align_corners, dim, ndims): def _scaled_indices(N, scale, align_corners, dim, ndims):
M = int(scale * N) M = int(scale * N)
if align_corners: if align_corners:
indices = mx.arange(M, dtype=mx.float32) * ((N - 1) / (M - 1)) indices = ((mx.arange(M, dtype=mx.float32) + 0.5) * (N / M) - 0.5).round()
else: else:
step = 1 / scale step = 1 / scale
start = ((M - 1) * step - N + 1) / 2 start = ((M - 1) * step - N + 1) / 2