pde_opt.numerics.functions.mixer_mlp

This module contains a Mixer MLP architecture for representing functions in PDEs.

Based on the code from https://docs.kidger.site/equinox/examples/score_based_diffusion/

Classes

`Mixer2d`(args, *kwargs)
`MixerBlock`(args, *kwargs)

class pde_opt.numerics.functions.mixer_mlp.MixerBlock(*args: Any, **kwargs: Any)[source]

__init__(num_patches, hidden_size, mix_patch_size, mix_hidden_size, *, key)[source]

patch_mixer: equinox.nn.MLP

hidden_mixer: equinox.nn.MLP

norm1: equinox.nn.LayerNorm

norm2: equinox.nn.LayerNorm

class pde_opt.numerics.functions.mixer_mlp.Mixer2d(*args: Any, **kwargs: Any)[source]

__init__(img_size, patch_size, hidden_size, mix_patch_size, mix_hidden_size, num_blocks, *, key)[source]

conv_in: equinox.nn.Conv2d

conv_out: equinox.nn.ConvTranspose2d

blocks: list

norm: equinox.nn.LayerNorm