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

Restyled pooling

Browse Source
This commit is contained in:
Vincent Amato
2025-08-11 23:15:38 -04:00
parent 652a143b64
commit c59b46a488
1 changed files with 0 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -451,12 +451,10 @@ class AdaptiveAvgPool1d(_AdaptivePool):
kernel_L = L // output_L kernel_L = L // output_L
if L % output_L == 0: if L % output_L == 0:
# Efficient path for exact division
new_shape = batch_dims + [output_L, kernel_L, C] new_shape = batch_dims + [output_L, kernel_L, C]
x_reshaped = x.reshape(new_shape) x_reshaped = x.reshape(new_shape)
return mx.mean(x_reshaped, axis=-2) return mx.mean(x_reshaped, axis=-2)
else: else:
# Manual indexing for non-exact division
stride_L = (L - kernel_L) // (output_L - 1) if output_L > 1 else 1 stride_L = (L - kernel_L) // (output_L - 1) if output_L > 1 else 1
values = [] values = []