monster/lato.2_mrp_mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Pass mode='max' explicitly to downsample

Browse Source 
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.
This commit is contained in:
John 2026-08-02 10:42:18 +10:00
parent 21c9de4bdd
commit ee86bdb4e4
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
lato_mlx/models/vvae.py
View File

@ -46,7 +46,10 @@ class DownResBlock(nn.Module):
self.block = SparseResBlock(channels, out_channels) self.block = SparseResBlock(channels, out_channels)
def __call__(self, x: SparseTensor) -> SparseTensor: 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): class SparseTransformerBase(nn.Module):