mlx 1.0.0

data/mlx/python/mlx/nn/layers/transformer.py

@@ -0,0 +1,354 @@
+# Copyright © 2023 Apple Inc.
+
+import math
+from typing import Any, Callable, Optional
+
+import mlx.core as mx
+from mlx.nn.layers.activations import relu
+from mlx.nn.layers.base import Module
+from mlx.nn.layers.dropout import Dropout
+from mlx.nn.layers.linear import Linear
+from mlx.nn.layers.normalization import LayerNorm
+from mlx.nn.utils import checkpoint
+
+
+class MultiHeadAttention(Module):
+    """Implements the scaled dot product attention with multiple heads.
+
+    Given inputs for queries, keys and values the ``MultiHeadAttention``
+    produces new values by aggregating information from the input values
+    according to the similarities of the input queries and keys.
+
+    All inputs as well as the output are linearly projected without biases by
+    default.
+
+    ``MultiHeadAttention`` also takes an optional additive attention mask that
+    should be broadcastable with ``(batch, num_heads, # queries, # keys)``. The
+    mask should have ``-inf`` or very large negative numbers at the positions
+    that should *not* be attended to.
+
+    Args:
+        dims (int): The model dimensions. This is also the default
+            value for the queries, keys, values, and the output.
+        num_heads (int): The number of attention heads to use.
+        query_input_dims (int, optional): The input dimensions of the queries.
+            Default: ``dims``.
+        key_input_dims (int, optional): The input dimensions of the keys.
+            Default: ``dims``.
+        value_input_dims (int, optional): The input dimensions of the values.
+            Default: ``key_input_dims``.
+        value_dims (int, optional): The dimensions of the values after the
+            projection. Default: ``dims``.
+        value_output_dims (int, optional): The dimensions the new values will
+            be projected to. Default: ``dims``.
+        bias (bool, optional): Whether or not to use a bias in the projections.
+            Default: ``False``.
+    """
+
+    def __init__(
+        self,
+        dims: int,
+        num_heads: int,
+        query_input_dims: Optional[int] = None,
+        key_input_dims: Optional[int] = None,
+        value_input_dims: Optional[int] = None,
+        value_dims: Optional[int] = None,
+        value_output_dims: Optional[int] = None,
+        bias: bool = False,
+    ):
+        super().__init__()
+
+        if (dims % num_heads) != 0:
+            raise ValueError(
+                "The input feature dimensions should be divisible by the "
+                f"number of heads ({dims} % {num_heads}) != 0"
+            )
+
+        query_input_dims = query_input_dims or dims
+        key_input_dims = key_input_dims or dims
+        value_input_dims = value_input_dims or key_input_dims
+        value_dims = value_dims or dims
+        value_output_dims = value_output_dims or dims
+
+        self.num_heads = num_heads
+        self.query_proj = Linear(query_input_dims, dims, bias=bias)
+        self.key_proj = Linear(key_input_dims, dims, bias=bias)
+        self.value_proj = Linear(value_input_dims, value_dims, bias=bias)
+        self.out_proj = Linear(value_dims, value_output_dims, bias=bias)
+
+    def __call__(self, queries, keys, values, mask=None):
+        queries = self.query_proj(queries)
+        keys = self.key_proj(keys)
+        values = self.value_proj(values)
+
+        num_heads = self.num_heads
+        queries = mx.unflatten(queries, -1, (num_heads, -1)).transpose(0, 2, 1, 3)
+        keys = mx.unflatten(keys, -1, (num_heads, -1)).transpose(0, 2, 1, 3)
+        values = mx.unflatten(values, -1, (num_heads, -1)).transpose(0, 2, 1, 3)
+        scale = math.sqrt(1 / queries.shape[-1])
+        output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).flatten(-2, -1)
+        return self.out_proj(output)
+
+    @staticmethod
+    def create_additive_causal_mask(N: int, dtype: mx.Dtype = mx.float32):
+        indices = mx.arange(N)
+        mask = indices[:, None] < indices[None]
+        mask = mask.astype(dtype) * mx.finfo(dtype).min
+        return mask
+
+
+class TransformerEncoderLayer(Module):
+    def __init__(
+        self,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: Callable[[Any], Any] = relu,
+        norm_first: bool = True,
+    ):
+        super().__init__()
+        mlp_dims = mlp_dims or dims * 4
+        self.attention = MultiHeadAttention(dims, num_heads)
+        self.ln1 = LayerNorm(dims)
+        self.ln2 = LayerNorm(dims)
+        self.linear1 = Linear(dims, mlp_dims)
+        self.linear2 = Linear(mlp_dims, dims)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+        self.activation = activation
+        self.norm_first = norm_first
+
+    def __call__(self, x, mask):
+        if self.norm_first:
+            y = self.ln1(x)
+            y = self.attention(y, y, y, mask)
+            y = self.dropout1(y)
+            x = x + y
+
+            y = self.ln2(x)
+            y = self.linear1(y)
+            y = self.activation(y)
+            y = self.dropout2(y)
+            y = self.linear2(y)
+            y = x + y
+
+        else:
+            y = self.attention(x, x, x, mask)
+            y = self.dropout1(y)
+            x = self.ln1(x + y)
+
+            y = self.linear1(x)
+            y = self.activation(y)
+            y = self.dropout2(y)
+            y = self.linear2(y)
+            y = self.ln2(x + y)
+
+        return y
+
+
+class TransformerEncoder(Module):
+    def __init__(
+        self,
+        num_layers: int,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation=relu,
+        norm_first: bool = True,
+        checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.layers = [
+            TransformerEncoderLayer(
+                dims, num_heads, mlp_dims, dropout, activation, norm_first
+            )
+            for i in range(num_layers)
+        ]
+        self.ln = LayerNorm(dims)
+        self.checkpoint = checkpoint
+
+    def __call__(self, x, mask):
+        for l in self.layers:
+            l = checkpoint(l) if self.checkpoint else l
+            x = l(x, mask)
+        return self.ln(x)
+
+
+class TransformerDecoderLayer(Module):
+    def __init__(
+        self,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: Callable[[Any], Any] = relu,
+        norm_first: bool = True,
+    ):
+        super().__init__()
+        mlp_dims = mlp_dims or dims * 4
+        self.self_attention = MultiHeadAttention(dims, num_heads)
+        self.cross_attention = MultiHeadAttention(dims, num_heads)
+        self.ln1 = LayerNorm(dims)
+        self.ln2 = LayerNorm(dims)
+        self.ln3 = LayerNorm(dims)
+        self.linear1 = Linear(dims, mlp_dims)
+        self.linear2 = Linear(mlp_dims, dims)
+        self.dropout1 = Dropout(dropout)
+        self.dropout2 = Dropout(dropout)
+        self.dropout3 = Dropout(dropout)
+        self.activation = activation
+        self.norm_first = norm_first
+
+    def __call__(self, x, memory, x_mask, memory_mask):
+        if self.norm_first:
+            y = self.ln1(x)
+            y = self.self_attention(y, y, y, x_mask)
+            y = self.dropout1(y)
+            x = x + y
+
+            y = self.ln2(x)
+            y = self.cross_attention(y, memory, memory, memory_mask)
+            y = self.dropout2(y)
+            x = x + y
+
+            y = self.ln3(x)
+            y = self.linear1(y)
+            y = self.activation(y)
+            y = self.dropout3(y)
+            y = self.linear2(y)
+            y = x + y
+
+        else:
+            y = self.self_attention(x, x, x, x_mask)
+            y = self.dropout1(y)
+            x = self.ln1(x + y)
+
+            y = self.cross_attention(y, memory, memory, memory_mask)
+            y = self.dropout2(y)
+            x = self.ln2(x + y)
+
+            y = self.linear1(x)
+            y = self.activation(y)
+            y = self.dropout3(y)
+            y = self.linear2(y)
+            y = self.ln3(x + y)
+
+        return y
+
+
+class TransformerDecoder(Module):
+    def __init__(
+        self,
+        num_layers: int,
+        dims: int,
+        num_heads: int,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation=relu,
+        norm_first: bool = True,
+        checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.layers = [
+            TransformerDecoderLayer(
+                dims, num_heads, mlp_dims, dropout, activation, norm_first
+            )
+            for i in range(num_layers)
+        ]
+        self.ln = LayerNorm(dims)
+        self.checkpoint = checkpoint
+
+    def __call__(self, x, memory, x_mask, memory_mask):
+        for l in self.layers:
+            l = checkpoint(l) if self.checkpoint else l
+            x = l(x, memory, x_mask, memory_mask)
+        return self.ln(x)
+
+
+class Transformer(Module):
+    """
+    Implements a standard Transformer model.
+
+    The implementation is based on `Attention Is All You Need
+    <https://arxiv.org/abs/1706.03762>`_.
+
+    The Transformer model contains an encoder and a decoder. The encoder
+    processes the input sequence and the decoder generates the output sequence.
+    The interaction between encoder and decoder happens through the attention
+    mechanism.
+
+    Args:
+        dims (int, optional): The number of expected features in the
+            encoder/decoder inputs. Default: ``512``.
+        num_heads (int, optional): The number of attention heads. Default:
+            ``8``.
+        num_encoder_layers (int, optional): The number of encoder layers in the
+            Transformer encoder. Default: ``6``.
+        num_decoder_layers (int, optional): The number of decoder layers in the
+            Transformer decoder. Default: ``6``.
+        mlp_dims (int, optional): The hidden dimension of the MLP block in each
+            Transformer layer. Defaults to ``4*dims`` if not provided. Default:
+            ``None``.
+        dropout (float, optional): The dropout value for the Transformer
+            encoder and decoder. Dropout is used after each attention layer and
+            the activation in the MLP layer. Default: ``0.0``.
+        activation (function, optional): the activation function for the MLP
+            hidden layer. Default: :func:`mlx.nn.relu`.
+        custom_encoder (nn.Module, optional): A custom encoder to replace the
+            standard Transformer encoder. Default: ``None``.
+        custom_decoder (nn.Module, optional): A custom decoder to replace the
+            standard Transformer decoder. Default: ``None``.
+        norm_first (bool, optional): if ``True``, encoder and decoder layers
+            will perform layer normalization before attention and MLP
+            operations, otherwise after. Default: ``True``.
+        checkpoint (bool, optional): if ``True`` perform gradient checkpointing
+            to reduce the memory usage at the expense of more computation.
+            Default: ``False``.
+    """
+
+    def __init__(
+        self,
+        dims: int = 512,
+        num_heads: int = 8,
+        num_encoder_layers: int = 6,
+        num_decoder_layers: int = 6,
+        mlp_dims: Optional[int] = None,
+        dropout: float = 0.0,
+        activation: Callable[[Any], Any] = relu,
+        custom_encoder: Optional[Any] = None,
+        custom_decoder: Optional[Any] = None,
+        norm_first: bool = True,
+        checkpoint: bool = False,
+    ):
+        super().__init__()
+
+        self.encoder = custom_encoder or TransformerEncoder(
+            num_encoder_layers,
+            dims,
+            num_heads,
+            mlp_dims,
+            dropout,
+            activation,
+            norm_first,
+            checkpoint,
+        )
+
+        self.decoder = custom_decoder or TransformerDecoder(
+            num_decoder_layers,
+            dims,
+            num_heads,
+            mlp_dims,
+            dropout,
+            activation,
+            norm_first,
+            checkpoint,
+        )
+
+    def __call__(self, src, tgt, src_mask, tgt_mask, memory_mask):
+        memory = self.encoder(src, src_mask)
+        return self.decoder(tgt, memory, tgt_mask, memory_mask)

data/mlx/python/mlx/nn/layers/upsample.py

@@ -0,0 +1,277 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import operator
+from functools import partial, reduce
+from itertools import product
+from typing import Callable, Literal, Tuple, Union
+
+import mlx.core as mx
+from mlx.nn.layers.base import Module
+
+
+def _scaled_indices(N, scale, align_corners, dim, ndims):
+    M = int(scale * N)
+    if align_corners:
+        indices = mx.arange(M, dtype=mx.float32) * ((N - 1) / (M - 1))
+    else:
+        step = 1 / scale
+        start = ((M - 1) * step - N + 1) / 2
+        indices = mx.arange(M, dtype=mx.float32) * step - start
+
+    shape = [1] * ndims
+    shape[dim] = -1
+
+    return indices.reshape(shape)
+
+
+def _nearest_indices(N, scale, dim, ndims):
+    M = int(scale * N)
+    indices = mx.arange(M, dtype=mx.float32)
+    if M > N:
+        indices = (indices + 0.5) * (N / M) - 0.5
+        indices = indices.round()
+    else:
+        indices = indices * (N / M)
+    shape = [1] * ndims
+    shape[dim] = -1
+    return indices.astype(mx.uint32).reshape(shape)
+
+
+def _linear_indices(N, scale, align_corners, dim, ndims):
+    indices = _scaled_indices(N, scale, align_corners, dim, ndims)
+    indices = mx.clip(indices, a_min=0, a_max=N - 1)
+    indices_l = mx.floor(indices)
+    indices_r = mx.ceil(indices)
+    weight = indices - indices_l
+    weight = mx.expand_dims(weight, -1)
+
+    return (
+        (indices_l.astype(mx.uint32), 1 - weight),
+        (indices_r.astype(mx.uint32), weight),
+    )
+
+
+def _cubic_indices(N, scale, align_corners, dim, ndims):
+    indices = _scaled_indices(N, scale, align_corners, dim, ndims)
+    indices_l1 = mx.floor(indices)
+    indices_r1 = mx.floor(indices + 1)
+    indices_l2 = indices_l1 - 1
+    indices_r2 = indices_r1 + 1
+
+    @partial(mx.compile, shapeless=True)
+    def _get_weight(ind, grid, dist):
+        # PyTorch uses -0.5 for antialiasing=true (compatibility with PIL)
+        # and uses -0.75 for antialiasing=false (compatibility with OpenCV)
+        a = -0.75
+        x = mx.abs(ind - grid)
+        if dist == 1:
+            weight = ((a + 2.0) * x - (a + 3.0)) * x * x + 1
+        else:
+            weight = (((x - 5) * x + 8) * x - 4) * a
+        return weight
+
+    weight_l1 = _get_weight(indices, indices_l1, dist=1)[..., None]
+    weight_r1 = _get_weight(indices, indices_r1, dist=1)[..., None]
+    weight_l2 = _get_weight(indices, indices_l2, dist=2)[..., None]
+    weight_r2 = _get_weight(indices, indices_r2, dist=2)[..., None]
+
+    # padding with border value
+    indices_l1 = mx.clip(indices_l1, a_min=0, a_max=N - 1)
+    indices_r1 = mx.clip(indices_r1, a_min=0, a_max=N - 1)
+    indices_l2 = mx.clip(indices_l2, a_min=0, a_max=N - 1)
+    indices_r2 = mx.clip(indices_r2, a_min=0, a_max=N - 1)
+
+    return (
+        (indices_l1.astype(mx.uint32), weight_l1),
+        (indices_r1.astype(mx.uint32), weight_r1),
+        (indices_l2.astype(mx.uint32), weight_l2),
+        (indices_r2.astype(mx.uint32), weight_r2),
+    )
+
+
+def upsample_nearest(x: mx.array, scale_factor: Tuple):
+    dims = x.ndim - 2
+    if dims != len(scale_factor):
+        raise ValueError("A scale needs to be provided for each spatial dimension")
+
+    # Integer scale_factors means we can simply expand-broadcast and reshape
+    if tuple(map(int, scale_factor)) == scale_factor:
+        shape = list(x.shape)
+        for d in range(dims):
+            shape.insert(2 + 2 * d, 1)
+        x = x.reshape(shape)
+        for d in range(dims):
+            shape[2 + 2 * d] = int(scale_factor[d])
+        x = mx.broadcast_to(x, shape)
+        for d in range(dims):
+            shape[d + 1] *= shape[d + 2]
+            shape.pop(d + 2)
+        x = x.reshape(shape)
+        return x
+
+    else:
+        B, *N, C = x.shape
+        indices = [slice(None)]
+        for i, (n, s) in enumerate(zip(N, scale_factor)):
+            indices.append(_nearest_indices(n, s, i, dims))
+        indices = tuple(indices)
+
+        return x[indices]
+
+
+def _interpolate(
+    x: mx.array, scale_factor: Tuple, indices_fn: Callable, align_corners: bool = False
+):
+    dims = x.ndim - 2
+    if dims != len(scale_factor):
+        raise ValueError("A scale needs to be provided for each spatial dimension")
+
+    B, *N, C = x.shape
+
+    # Compute the sampling grid
+    indices = []
+    for i, (n, s) in enumerate(zip(N, scale_factor)):
+        indices.append(indices_fn(n, s, align_corners, i, dims))
+
+    # Sample and compute the weights
+    samples = []
+    weights = []
+    for idx_weight in product(*indices):
+        idx, weight = zip(*idx_weight)
+        samples.append(x[(slice(None),) + idx])
+        weights.append(reduce(operator.mul, weight))
+
+    # Interpolate
+    return sum(wi * xi for wi, xi in zip(weights, samples))
+
+
+def upsample_linear(x: mx.array, scale_factor: Tuple, align_corners: bool = False):
+    return _interpolate(
+        x=x,
+        scale_factor=scale_factor,
+        indices_fn=_linear_indices,
+        align_corners=align_corners,
+    )
+
+
+def upsample_cubic(x: mx.array, scale_factor: Tuple, align_corners: bool = False):
+    return _interpolate(
+        x=x,
+        scale_factor=scale_factor,
+        indices_fn=_cubic_indices,
+        align_corners=align_corners,
+    )
+
+
+class Upsample(Module):
+    r"""Upsample the input signal spatially.
+
+    The spatial dimensions are by convention dimensions ``1`` to ``x.ndim -
+    2``. The first is the batch dimension and the last is the feature
+    dimension.
+
+    For example, an audio signal would be 3D with 1 spatial dimension, an image
+    4D with 2 and so on and so forth.
+
+    There are three upsampling algorithms implemented nearest neighbor upsampling,
+    linear interpolation, and cubic interpolation. All can be applied to any number
+    of spatial dimensions. The linear interpolation will be bilinear, trilinear etc
+    when applied to more than one spatial dimension. And cubic interpolation will be
+    bicubic when there are 2 spatial dimensions.
+
+    .. note::
+       When using one of the linear or cubic interpolation modes the ``align_corners``
+       argument changes how the corners are treated in the input image. If
+       ``align_corners=True`` then the top and left edge of the input and
+       output will be matching as will the bottom right edge.
+
+    Parameters:
+        scale_factor (float or tuple): The multiplier for the spatial size.
+            If a ``float`` is provided, it is the multiplier for all spatial dimensions.
+            Otherwise, the number of scale factors provided must match the
+            number of spatial dimensions.
+        mode (str, optional): The upsampling algorithm, either ``"nearest"``,
+            ``"linear"`` or ``"cubic"``. Default: ``"nearest"``.
+        align_corners (bool, optional): Changes the way the corners are treated
+            during ``"linear"`` and ``"cubic"`` upsampling.  See the note above and the
+            examples below for more details.  Default: ``False``.
+
+    Examples:
+        >>> import mlx.core as mx
+        >>> import mlx.nn as nn
+        >>> x = mx.arange(1, 5).reshape((1, 2, 2, 1))
+        >>> x
+        array([[[[1],
+                 [2]],
+                [[3],
+                 [4]]]], dtype=int32)
+        >>> n = nn.Upsample(scale_factor=2, mode='nearest')
+        >>> n(x).squeeze()
+        array([[1, 1, 2, 2],
+               [1, 1, 2, 2],
+               [3, 3, 4, 4],
+               [3, 3, 4, 4]], dtype=int32)
+        >>> b = nn.Upsample(scale_factor=2, mode='linear')
+        >>> b(x).squeeze()
+        array([[1, 1.25, 1.75, 2],
+               [1.5, 1.75, 2.25, 2.5],
+               [2.5, 2.75, 3.25, 3.5],
+               [3, 3.25, 3.75, 4]], dtype=float32)
+        >>> b = nn.Upsample(scale_factor=2, mode='linear', align_corners=True)
+        >>> b(x).squeeze()
+        array([[1, 1.33333, 1.66667, 2],
+               [1.66667, 2, 2.33333, 2.66667],
+               [2.33333, 2.66667, 3, 3.33333],
+               [3, 3.33333, 3.66667, 4]], dtype=float32)
+    """
+
+    def __init__(
+        self,
+        scale_factor: Union[float, Tuple],
+        mode: Literal["nearest", "linear", "cubic"] = "nearest",
+        align_corners: bool = False,
+    ):
+        super().__init__()
+        if mode not in ["nearest", "linear", "cubic"]:
+            raise ValueError(f"[Upsample] Got unsupported upsampling algorithm: {mode}")
+        if isinstance(scale_factor, (list, tuple)):
+            self.scale_factor = tuple(map(float, scale_factor))
+        else:
+            self.scale_factor = float(scale_factor)
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def _extra_repr(self) -> str:
+        return (
+            f"scale_factor={self.scale_factor}, mode={self.mode!r}, "
+            f"align_corners={self.align_corners}"
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        dims = x.ndim - 2
+        if dims <= 0:
+            raise ValueError(
+                f"[Upsample] The input should have at least 1 spatial "
+                f"dimension which means it should be at least 3D but "
+                f"{x.ndim}D was provided"
+            )
+
+        scale_factor = self.scale_factor
+        if isinstance(scale_factor, tuple):
+            if len(scale_factor) != dims:
+                raise ValueError(
+                    f"[Upsample] One scale per spatial dimension is required but "
+                    f"scale_factor={scale_factor} and the number of spatial "
+                    f"dimensions were {dims}"
+                )
+        else:
+            scale_factor = (scale_factor,) * dims
+
+        if self.mode == "nearest":
+            return upsample_nearest(x, scale_factor)
+        elif self.mode == "linear":
+            return upsample_linear(x, scale_factor, self.align_corners)
+        elif self.mode == "cubic":
+            return upsample_cubic(x, scale_factor, self.align_corners)
+        else:
+            raise Exception(f"Unknown interpolation mode: {self.mode}")