"""MLX implementation of the TRELLIS-lineage sparse module.

Microsoft's TRELLIS.2 sparse stack has been inherited, near-verbatim, by a growing
family of 3D generation models — LATO.2 and TencentARC's Pixal3D among them. All of
them hard-require spconv or torchsparse, neither of which has a Metal build, and that
single dependency is what keeps the whole lineage off Apple Silicon.

The blocker is one operation: submanifold 3x3x3 convolution. Every SparseConv3d in
these models is constructed `stride=1, padding=None`, which spconv dispatches to
SubMConv3d. Implement that in MLX and the rest is ordinary linear/norm/attention work.

Packaged separately from any one model so each port depends on a tested core rather
than vendoring its own copy.
"""

from .conv import SubMConv3d, build_indice_map
from .dit import (
    ProjectAttention,
    DiTAttention,
    ModulatedTransformerCrossBlock,
    MultiHeadRMSNorm,
    TimestepEmbedder,
    apply_rope,
    rope_phases_from_coords,
    SparseDiTAttention,
    SparseProjectAttention,
    ModulatedSparseTransformerCrossBlock,
)
from .tensor import (
    SparseTensor,
    VarLenTensor,
    channel2spatial,
    downsample,
    spatial2channel,
    subdivide,
    upsample,
)
from .ops import (
    LayerNorm32,
    SparseFeedForwardNet,
    SparseGELU,
    SparseGroupNorm32,
    SparseLinear,
    SparseMultiHeadAttention,
    SparseResBlock,
    SparseSiLU,
    SparseTransformerBlock,
    SparseTransformerCrossBlock,
    SparseConvNeXtBlock3d,
    SparseResBlockC2S3d,
)

__all__ = [
    "SparseTensor", "VarLenTensor", "subdivide", "downsample", "upsample",
    "spatial2channel", "channel2spatial",
    "SubMConv3d", "build_indice_map",
    "SparseLinear", "LayerNorm32", "SparseGroupNorm32", "SparseSiLU", "SparseGELU",
    "SparseResBlock", "SparseFeedForwardNet", "SparseMultiHeadAttention",
    "SparseTransformerBlock", "SparseTransformerCrossBlock",
    "SparseConvNeXtBlock3d", "SparseResBlockC2S3d",
    # DiT / flow-model pieces
    "MultiHeadRMSNorm", "apply_rope", "rope_phases_from_coords", "TimestepEmbedder", "DiTAttention",
    "ModulatedTransformerCrossBlock", "ProjectAttention",
    "SparseDiTAttention", "SparseProjectAttention",
    "ModulatedSparseTransformerCrossBlock",
]
__version__ = "0.1.0"