Files
mlx-examples/normalizing_flow/distributions.py