mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 01:17:28 +08:00
19b6167d81
* Implement normalizing flow Real NVP example * Add requirements and basic usage to normalizing flow example * Minor changes to README in normalizing flow example * Remove trailing commas in function arguments for unified formatting in flows example * Fix minor typos, add some annotations * format + nits in README * readme fix * mov, minor changes in main, copywright * remove debug * fix * Simplified class structure in distributions; better code re-use in bijectors * Remove rogue space * change name again * nits --------- Co-authored-by: Awni Hannun <awni@apple.com>
61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
# Copyright © 2023-2024 Apple Inc.
|
|
|
|
from typing import Tuple
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
|
|
class Bijector:
|
|
def forward_and_log_det(self, x: mx.array) -> Tuple[mx.array, mx.array]:
|
|
raise NotImplementedError
|
|
|
|
def inverse_and_log_det(self, y: mx.array) -> Tuple[mx.array, mx.array]:
|
|
raise NotImplementedError
|
|
|
|
|
|
class AffineBijector(Bijector):
|
|
def __init__(self, shift_and_log_scale: mx.array):
|
|
self.shift_and_log_scale = shift_and_log_scale
|
|
|
|
def forward_and_log_det(self, x: mx.array):
|
|
shift, log_scale = mx.split(self.shift_and_log_scale, 2, axis=-1)
|
|
y = x * mx.exp(log_scale) + shift
|
|
log_det = log_scale
|
|
return y, log_det
|
|
|
|
def inverse_and_log_det(self, y: mx.array):
|
|
shift, log_scale = mx.split(self.shift_and_log_scale, 2, axis=-1)
|
|
x = (y - shift) * mx.exp(-log_scale)
|
|
log_det = -log_scale
|
|
return x, log_det
|
|
|
|
|
|
class MaskedCoupling(Bijector):
|
|
def __init__(self, mask: mx.array, conditioner: nn.Module, bijector: Bijector):
|
|
"""Coupling layer with masking and conditioner."""
|
|
self.mask = mask
|
|
self.conditioner = conditioner
|
|
self.bijector = bijector
|
|
|
|
def apply_mask(self, x: mx.array, func: callable):
|
|
"""Transforms masked indices of `x` conditioned on unmasked indices using `func`."""
|
|
x_masked = mx.where(self.mask, 0.0, x)
|
|
bijector_params = self.conditioner(x_masked)
|
|
y, log_det = func(bijector_params)
|
|
log_det = mx.where(self.mask, log_det, 0.0)
|
|
y = mx.where(self.mask, y, x)
|
|
return y, mx.sum(log_det, axis=-1)
|
|
|
|
def forward_and_log_det(self, x: mx.array):
|
|
"""Transforms masked indices of `x` conditioned on unmasked indices using bijector."""
|
|
return self.apply_mask(
|
|
x, lambda params: self.bijector(params).forward_and_log_det(x)
|
|
)
|
|
|
|
def inverse_and_log_det(self, y: mx.array):
|
|
"""Transforms masked indices of `y` conditioned on unmasked indices using bijector."""
|
|
return self.apply_mask(
|
|
y, lambda params: self.bijector(params).inverse_and_log_det(y)
|
|
)
|