mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 01:17:28 +08:00
32 lines
850 B
Python
32 lines
850 B
Python
# Copyright © 2023-2024 Apple Inc.
|
|
|
|
import math
|
|
from typing import Optional, Tuple, Union
|
|
|
|
import mlx.core as mx
|
|
|
|
|
|
class Normal:
|
|
def __init__(self, mu: mx.array, sigma: mx.array):
|
|
super().__init__()
|
|
self.mu = mu
|
|
self.sigma = sigma
|
|
|
|
def sample(
|
|
self, sample_shape: Union[int, Tuple[int, ...]], key: Optional[mx.array] = None
|
|
):
|
|
return mx.random.normal(sample_shape, key=key) * self.sigma + self.mu
|
|
|
|
def log_prob(self, x: mx.array):
|
|
return (
|
|
-0.5 * math.log(2 * math.pi)
|
|
- mx.log(self.sigma)
|
|
- 0.5 * ((x - self.mu) / self.sigma) ** 2
|
|
)
|
|
|
|
def sample_and_log_prob(
|
|
self, sample_shape: Union[int, Tuple[int, ...]], key: Optional[mx.array] = None
|
|
):
|
|
x = self.sample(sample_shape, key=key)
|
|
return x, self.log_prob(x)
|