From ee86bdb4e40a2ecdb134c094c7d40e1e8e35cfda Mon Sep 17 00:00:00 2001 From: John Date: Sun, 2 Aug 2026 10:42:18 +1000 Subject: [PATCH] Pass mode='max' explicitly to downsample LATO.2 hardcodes reduce='amax' but Pixal3D's version of the same op defaults to 'mean'. Both now share one implementation, so relying on its default would silently change every downsampled feature if that default ever moved. --- lato_mlx/models/vvae.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/lato_mlx/models/vvae.py b/lato_mlx/models/vvae.py index 5065786..02f2335 100644 --- a/lato_mlx/models/vvae.py +++ b/lato_mlx/models/vvae.py @@ -46,7 +46,10 @@ class DownResBlock(nn.Module): self.block = SparseResBlock(channels, out_channels) def __call__(self, x: SparseTensor) -> SparseTensor: - return self.block(downsample(x, 2)) + # mode="max" explicitly: LATO.2 hardcodes reduce="amax", but Pixal3D's version of + # the same op defaults to "mean". Relying on the shared default would silently + # change every downsampled feature the day that default moves. + return self.block(downsample(x, 2, mode="max")) class SparseTransformerBase(nn.Module):