|
|
|
|
|
|
|
|
#include <nanobind/nanobind.h> |
|
|
#include <nanobind/stl/optional.h> |
|
|
#include <nanobind/stl/variant.h> |
|
|
#include <nanobind/stl/vector.h> |
|
|
|
|
|
#include <chrono> |
|
|
|
|
|
#include "mlx/ops.h" |
|
|
#include "mlx/random.h" |
|
|
#include "python/src/small_vector.h" |
|
|
#include "python/src/utils.h" |
|
|
|
|
|
namespace mx = mlx::core; |
|
|
namespace nb = nanobind; |
|
|
using namespace nb::literals; |
|
|
|
|
|
class PyKeySequence { |
|
|
public: |
|
|
explicit PyKeySequence(uint64_t seed) { |
|
|
state_.append(mx::random::key(seed)); |
|
|
} |
|
|
|
|
|
void seed(uint64_t seed) { |
|
|
state_[0] = mx::random::key(seed); |
|
|
} |
|
|
|
|
|
mx::array next() { |
|
|
auto out = mx::random::split(nb::cast<mx::array>(state_[0])); |
|
|
state_[0] = out.first; |
|
|
return out.second; |
|
|
} |
|
|
|
|
|
nb::list state() { |
|
|
return state_; |
|
|
} |
|
|
|
|
|
void release() { |
|
|
nb::gil_scoped_acquire gil; |
|
|
state_.release().dec_ref(); |
|
|
} |
|
|
|
|
|
private: |
|
|
nb::list state_; |
|
|
}; |
|
|
|
|
|
PyKeySequence& default_key() { |
|
|
auto get_current_time_seed = []() { |
|
|
auto now = std::chrono::system_clock::now(); |
|
|
return std::chrono::duration_cast<std::chrono::milliseconds>( |
|
|
now.time_since_epoch()) |
|
|
.count(); |
|
|
}; |
|
|
static PyKeySequence ks(get_current_time_seed()); |
|
|
return ks; |
|
|
} |
|
|
|
|
|
void init_random(nb::module_& parent_module) { |
|
|
auto m = parent_module.def_submodule( |
|
|
"random", |
|
|
"mlx.core.random: functionality related to random number generation"); |
|
|
|
|
|
m.attr("state") = default_key().state(); |
|
|
m.def( |
|
|
"seed", |
|
|
[](uint64_t seed) { default_key().seed(seed); }, |
|
|
"seed"_a, |
|
|
R"pbdoc( |
|
|
Seed the global PRNG. |
|
|
|
|
|
Args: |
|
|
seed (int): Seed for the global PRNG. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"key", |
|
|
&mx::random::key, |
|
|
"seed"_a, |
|
|
R"pbdoc( |
|
|
Get a PRNG key from a seed. |
|
|
|
|
|
Args: |
|
|
seed (int): Seed for the PRNG. |
|
|
|
|
|
Returns: |
|
|
array: The PRNG key array. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"split", |
|
|
nb::overload_cast<const mx::array&, int, mx::StreamOrDevice>( |
|
|
&mx::random::split), |
|
|
"key"_a, |
|
|
"num"_a = 2, |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def split(key: array, num: int = 2, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Split a PRNG key into sub keys. |
|
|
|
|
|
Args: |
|
|
key (array): Input key to split. |
|
|
num (int, optional): Number of sub keys. Default: ``2``. |
|
|
|
|
|
Returns: |
|
|
array: The array of sub keys with ``num`` as its first dimension. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"uniform", |
|
|
[](const ScalarOrArray& low, |
|
|
const ScalarOrArray& high, |
|
|
const mx::Shape& shape, |
|
|
std::optional<mx::Dtype> type, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
return mx::random::uniform( |
|
|
to_array(low), |
|
|
to_array(high), |
|
|
shape, |
|
|
type.value_or(mx::float32), |
|
|
key, |
|
|
s); |
|
|
}, |
|
|
"low"_a = 0, |
|
|
"high"_a = 1, |
|
|
"shape"_a = mx::Shape{}, |
|
|
"dtype"_a.none() = mx::float32, |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def uniform(low: Union[scalar, array] = 0, high: Union[scalar, array] = 1, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Generate uniformly distributed random numbers. |
|
|
|
|
|
The values are sampled uniformly in the half-open interval ``[low, high)``. |
|
|
The lower and upper bound can be scalars or arrays and must be |
|
|
broadcastable to ``shape``. |
|
|
|
|
|
Args: |
|
|
low (scalar or array, optional): Lower bound of the distribution. |
|
|
Default: ``0``. |
|
|
high (scalar or array, optional): Upper bound of the distribution. |
|
|
Default: ``1``. |
|
|
shape (list(int), optional): Shape of the output. Default:``()``. |
|
|
dtype (Dtype, optional): Type of the output. Default: ``float32``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The output array random values. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"normal", |
|
|
[](const mx::Shape& shape, |
|
|
std::optional<mx::Dtype> type, |
|
|
const std::optional<ScalarOrArray>& loc_, |
|
|
const std::optional<ScalarOrArray>& scale_, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto dtype = type.value_or(mx::float32); |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
auto loc = |
|
|
loc_ ? std::make_optional(to_array(*loc_, dtype)) : std::nullopt; |
|
|
auto scale = scale_ ? std::make_optional(to_array(*scale_, dtype)) |
|
|
: std::nullopt; |
|
|
return mx::random::normal(shape, dtype, loc, scale, key, s); |
|
|
}, |
|
|
"shape"_a = mx::Shape{}, |
|
|
"dtype"_a.none() = mx::float32, |
|
|
"loc"_a = nb::none(), |
|
|
"scale"_a = nb::none(), |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def normal(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: Union[scalar, array, None] = None, scale: Union[scalar, array, None] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Generate normally distributed random numbers. |
|
|
|
|
|
If ``loc`` and ``scale`` are not provided the "standard" normal |
|
|
distribution is used. That means $x \sim \mathcal{N}(0, 1)$ for |
|
|
real numbers and $\text{Re}(x),\text{Im}(x) \sim \mathcal{N}(0, |
|
|
\frac{1}{2})$ for complex numbers. |
|
|
|
|
|
Args: |
|
|
shape (list(int), optional): Shape of the output. Default: ``()``. |
|
|
dtype (Dtype, optional): Type of the output. Default: ``float32``. |
|
|
loc (scalar or array, optional): Mean of the distribution. |
|
|
Default: ``None``. |
|
|
scale (scalar or array, optional): Standard deviation of the |
|
|
distribution. Default: ``None``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The output array of random values. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"multivariate_normal", |
|
|
[](const mx::array& mean, |
|
|
const mx::array& cov, |
|
|
const mx::Shape& shape, |
|
|
std::optional<mx::Dtype> type, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
return mx::random::multivariate_normal( |
|
|
mean, cov, shape, type.value_or(mx::float32), key, s); |
|
|
}, |
|
|
"mean"_a, |
|
|
"cov"_a, |
|
|
"shape"_a = mx::Shape{}, |
|
|
"dtype"_a.none() = mx::float32, |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def multivariate_normal(mean: array, cov: array, shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Generate jointly-normal random samples given a mean and covariance. |
|
|
|
|
|
The matrix ``cov`` must be positive semi-definite. The behavior is |
|
|
undefined if it is not. The only supported ``dtype`` is ``float32``. |
|
|
|
|
|
Args: |
|
|
mean (array): array of shape ``(..., n)``, the mean of the |
|
|
distribution. |
|
|
cov (array): array of shape ``(..., n, n)``, the covariance |
|
|
matrix of the distribution. The batch shape ``...`` must be |
|
|
broadcast-compatible with that of ``mean``. |
|
|
shape (list(int), optional): The output shape must be |
|
|
broadcast-compatible with ``mean.shape[:-1]`` and ``cov.shape[:-2]``. |
|
|
If empty, the result shape is determined by broadcasting the batch |
|
|
shapes of ``mean`` and ``cov``. Default: ``[]``. |
|
|
dtype (Dtype, optional): The output type. Default: ``float32``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The output array of random values. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"randint", |
|
|
[](const ScalarOrArray& low, |
|
|
const ScalarOrArray& high, |
|
|
const mx::Shape& shape, |
|
|
std::optional<mx::Dtype> type, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
return mx::random::randint( |
|
|
to_array(low), |
|
|
to_array(high), |
|
|
shape, |
|
|
type.value_or(mx::int32), |
|
|
key, |
|
|
s); |
|
|
}, |
|
|
"low"_a, |
|
|
"high"_a, |
|
|
"shape"_a = mx::Shape{}, |
|
|
"dtype"_a.none() = mx::int32, |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def randint(low: Union[scalar, array], high: Union[scalar, array], shape: Sequence[int] = [], dtype: Optional[Dtype] = int32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Generate random integers from the given interval. |
|
|
|
|
|
The values are sampled with equal probability from the integers in |
|
|
half-open interval ``[low, high)``. The lower and upper bound can be |
|
|
scalars or arrays and must be broadcastable to ``shape``. |
|
|
|
|
|
Args: |
|
|
low (scalar or array): Lower bound of the interval. |
|
|
high (scalar or array): Upper bound of the interval. |
|
|
shape (list(int), optional): Shape of the output. Default: ``()``. |
|
|
dtype (Dtype, optional): Type of the output. Default: ``int32``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The array of random integers. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"bernoulli", |
|
|
[](const ScalarOrArray& p_, |
|
|
const std::optional<mx::Shape> shape, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
auto p = to_array(p_); |
|
|
if (shape.has_value()) { |
|
|
return mx::random::bernoulli(p, shape.value(), key, s); |
|
|
} else { |
|
|
return mx::random::bernoulli(p, key, s); |
|
|
} |
|
|
}, |
|
|
"p"_a = 0.5, |
|
|
"shape"_a = nb::none(), |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def bernoulli(p: Union[scalar, array] = 0.5, shape: Optional[Sequence[int]] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Generate Bernoulli random values. |
|
|
|
|
|
The values are sampled from the bernoulli distribution with parameter |
|
|
``p``. The parameter ``p`` can be a :obj:`float` or :obj:`array` and |
|
|
must be broadcastable to ``shape``. |
|
|
|
|
|
Args: |
|
|
p (float or array, optional): Parameter of the Bernoulli |
|
|
distribution. Default: ``0.5``. |
|
|
shape (list(int), optional): Shape of the output. |
|
|
Default: ``p.shape``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The array of random integers. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"truncated_normal", |
|
|
[](const ScalarOrArray& lower_, |
|
|
const ScalarOrArray& upper_, |
|
|
const std::optional<mx::Shape> shape_, |
|
|
std::optional<mx::Dtype> type, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
auto lower = to_array(lower_); |
|
|
auto upper = to_array(upper_); |
|
|
auto t = type.value_or(mx::float32); |
|
|
if (shape_.has_value()) { |
|
|
return mx::random::truncated_normal( |
|
|
lower, upper, shape_.value(), t, key, s); |
|
|
} else { |
|
|
return mx::random::truncated_normal(lower, upper, t, key, s); |
|
|
} |
|
|
}, |
|
|
"lower"_a, |
|
|
"upper"_a, |
|
|
"shape"_a = nb::none(), |
|
|
"dtype"_a.none() = mx::float32, |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def truncated_normal(lower: Union[scalar, array], upper: Union[scalar, array], shape: Optional[Sequence[int]] = None, dtype: Optional[Dtype] = float32, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Generate values from a truncated normal distribution. |
|
|
|
|
|
The values are sampled from the truncated normal distribution |
|
|
on the domain ``(lower, upper)``. The bounds ``lower`` and ``upper`` |
|
|
can be scalars or arrays and must be broadcastable to ``shape``. |
|
|
|
|
|
Args: |
|
|
lower (scalar or array): Lower bound of the domain. |
|
|
upper (scalar or array): Upper bound of the domain. |
|
|
shape (list(int), optional): The shape of the output. |
|
|
Default:``()``. |
|
|
dtype (Dtype, optional): The data type of the output. |
|
|
Default: ``float32``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The output array of random values. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"gumbel", |
|
|
[](const mx::Shape& shape, |
|
|
std::optional<mx::Dtype> type, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
return mx::random::gumbel(shape, type.value_or(mx::float32), key, s); |
|
|
}, |
|
|
"shape"_a = mx::Shape{}, |
|
|
"dtype"_a.none() = mx::float32, |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def gumbel(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, key: Union[None, Stream, Device] = None, stream: Optional[array] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Sample from the standard Gumbel distribution. |
|
|
|
|
|
The values are sampled from a standard Gumbel distribution |
|
|
which CDF ``exp(-exp(-x))``. |
|
|
|
|
|
Args: |
|
|
shape (list(int)): The shape of the output. |
|
|
dtype (Dtype, optional): The data type of the output. |
|
|
Default: ``float32``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: |
|
|
The :class:`array` with shape ``shape`` and distributed according |
|
|
to the Gumbel distribution. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"categorical", |
|
|
[](const mx::array& logits, |
|
|
int axis, |
|
|
const std::optional<mx::Shape> shape, |
|
|
const std::optional<int> num_samples, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
if (shape.has_value() && num_samples.has_value()) { |
|
|
throw std::invalid_argument( |
|
|
"[categorical] At most one of shape or num_samples can be specified."); |
|
|
} else if (shape.has_value()) { |
|
|
return mx::random::categorical(logits, axis, shape.value(), key, s); |
|
|
} else if (num_samples.has_value()) { |
|
|
return mx::random::categorical( |
|
|
logits, axis, num_samples.value(), key, s); |
|
|
} else { |
|
|
return mx::random::categorical(logits, axis, key, s); |
|
|
} |
|
|
}, |
|
|
"logits"_a, |
|
|
"axis"_a = -1, |
|
|
"shape"_a = nb::none(), |
|
|
"num_samples"_a = nb::none(), |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def categorical(logits: array, axis: int = -1, shape: Optional[Sequence[int]] = None, num_samples: Optional[int] = None, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Sample from a categorical distribution. |
|
|
|
|
|
The values are sampled from the categorical distribution specified by |
|
|
the unnormalized values in ``logits``. Note, at most one of ``shape`` |
|
|
or ``num_samples`` can be specified. If both are ``None``, the output |
|
|
has the same shape as ``logits`` with the ``axis`` dimension removed. |
|
|
|
|
|
Args: |
|
|
logits (array): The *unnormalized* categorical distribution(s). |
|
|
axis (int, optional): The axis which specifies the distribution. |
|
|
Default: ``-1``. |
|
|
shape (list(int), optional): The shape of the output. This must |
|
|
be broadcast compatible with ``logits.shape`` with the ``axis`` |
|
|
dimension removed. Default: ``None`` |
|
|
num_samples (int, optional): The number of samples to draw from each |
|
|
of the categorical distributions in ``logits``. The output will have |
|
|
``num_samples`` in the last dimension. Default: ``None``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The ``shape``-sized output array with type ``uint32``. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"laplace", |
|
|
[](const mx::Shape& shape, |
|
|
std::optional<mx::Dtype> type, |
|
|
float loc, |
|
|
float scale, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
return mx::random::laplace( |
|
|
shape, type.value_or(mx::float32), loc, scale, key, s); |
|
|
}, |
|
|
"shape"_a = mx::Shape{}, |
|
|
"dtype"_a.none() = mx::float32, |
|
|
"loc"_a = 0.0, |
|
|
"scale"_a = 1.0, |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def laplace(shape: Sequence[int] = [], dtype: Optional[Dtype] = float32, loc: float = 0.0, scale: float = 1.0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Sample numbers from a Laplace distribution. |
|
|
|
|
|
Args: |
|
|
shape (list(int), optional): Shape of the output. Default: ``()``. |
|
|
dtype (Dtype, optional): Type of the output. Default: ``float32``. |
|
|
loc (float, optional): Mean of the distribution. Default: ``0.0``. |
|
|
scale (float, optional): The scale "b" of the Laplace distribution. |
|
|
Default:``1.0``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: The output array of random values. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"permutation", |
|
|
[](const std::variant<nb::int_, mx::array>& x, |
|
|
int axis, |
|
|
const std::optional<mx::array>& key_, |
|
|
mx::StreamOrDevice s) { |
|
|
auto key = key_ ? key_.value() : default_key().next(); |
|
|
if (auto pv = std::get_if<nb::int_>(&x); pv) { |
|
|
return mx::random::permutation(nb::cast<int>(*pv), key, s); |
|
|
} else { |
|
|
return mx::random::permutation(std::get<mx::array>(x), axis, key, s); |
|
|
} |
|
|
}, |
|
|
"x"_a, |
|
|
"axis"_a = 0, |
|
|
"key"_a = nb::none(), |
|
|
"stream"_a = nb::none(), |
|
|
nb::sig( |
|
|
"def permutation(x: Union[int, array], axis: int = 0, key: Optional[array] = None, stream: Union[None, Stream, Device] = None) -> array"), |
|
|
R"pbdoc( |
|
|
Generate a random permutation or permute the entries of an array. |
|
|
|
|
|
Args: |
|
|
x (int or array, optional): If an integer is provided a random |
|
|
permtuation of ``mx.arange(x)`` is returned. Otherwise the entries |
|
|
of ``x`` along the given axis are randomly permuted. |
|
|
axis (int, optional): The axis to permute along. Default: ``0``. |
|
|
key (array, optional): A PRNG key. Default: ``None``. |
|
|
|
|
|
Returns: |
|
|
array: |
|
|
The generated random permutation or randomly permuted input array. |
|
|
)pbdoc"); |
|
|
|
|
|
auto atexit = nb::module_::import_("atexit"); |
|
|
atexit.attr("register")(nb::cpp_function([]() { default_key().release(); })); |
|
|
} |
|
|
|
