← back to Exo

.typings/mflux/models/flux/model/flux_vae/common/resnet_block_2d.pyi

22 lines

"""
This type stub file was generated by pyright.
"""

import mlx.core as mx
from mlx import nn

class ResnetBlock2D(nn.Module):
    def __init__(
        self,
        norm1: int,
        conv1_in: int,
        conv1_out: int,
        norm2: int,
        conv2_in: int,
        conv2_out: int,
        conv_shortcut_in: int | None = ...,
        conv_shortcut_out: int | None = ...,
        is_conv_shortcut: bool = ...,
    ) -> None: ...
    def __call__(self, input_array: mx.array) -> mx.array: ...