zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Awni Hannun
2025-07-15 13:12:45 -07:00
parent 7fe6f03a5b
commit 2d7e412298
2 changed files with 39 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
mlx/random.cpp
View File

@@ -82,6 +82,16 @@ array split(const array& key, int num, StreamOrDevice s /* = {} */) {
return bits({num, 2}, 4, key, s);
}
// Get the next representable value below 1.0 for half precision
// floating point types (fp16, bf16)
template <typename T>
T below_one() {
T f = T(1.0);
uint16_t* m = (uint16_t*)&f;
*m -= 1;
return f;
}
array uniform(
const array& low,
const array& high,
@@ -106,7 +116,23 @@ array uniform(
<< " from broadcasted shape " << out_shape << ".";
throw std::invalid_argument(msg.str());
}
auto upper = array(std::nextafter(1.0f, 0.0f), float32);
// Get random values between [0, nextafter(1.0, 0.0)] since samples must
// be in [low, high)
auto get_upper = [&dtype]() {
switch (dtype) {
case float32:
return array(std::nextafter(1.0f, 0.0f), float32);
case float16:
return array(below_one<float16_t>(), float32);
case bfloat16:
return array(below_one<bfloat16_t>(), float32);
default:
throw std::runtime_error("[uniform] Unsupported type.");
}
};
auto upper = get_upper();
auto maxval = array(std::numeric_limits<uint32_t>::max(), float32);
auto out = bits(shape, size_of(float32), key, stream);
out = divide(out, maxval, stream);
@@ -154,6 +180,10 @@ array normal(
StreamOrDevice s /* = {} */) {
if (dtype == complex64) {
return complex_normal(shape, loc, scale, key, s);
} else if (!issubdtype(dtype, floating)) {
throw std::invalid_argument(
"[normal] Can only generate uniform numbers with "
"floating point type.");
}
auto stream = to_stream(s);
@@ -417,6 +447,12 @@ array laplace(
const float scale /* = 1.0 */,
const std::optional<array>& key /*= nullopt */,
StreamOrDevice s /* = {} */) {
if (!issubdtype(dtype, floating)) {
throw std::invalid_argument(
"[laplace] Can only generate uniform numbers with real"
"floating point type.");
}
auto stream = to_stream(s);
auto low = array(std::nextafter(-1.0f, 0.0f), float32);
auto high = array(1.0f, float32);

4
tests/random_tests.cpp
View File

@@ -350,7 +350,7 @@ TEST_CASE("test random uniform") {
// Check float16
{
auto key = random::key(0);
auto out = random::uniform({100}, float16, key);
auto out = random::uniform({1000}, float16, key);
CHECK_EQ(out.dtype(), float16);
CHECK(all(less(out, array(1.0f))).item<bool>());
CHECK(all(greater_equal(out, array(0.0f))).item<bool>());
@@ -360,7 +360,7 @@ TEST_CASE("test random uniform") {
{
auto key = random::key(0);
auto out = random::uniform({100}, bfloat16, key);
auto out = random::uniform({1000}, bfloat16, key);
CHECK_EQ(out.dtype(), bfloat16);
CHECK(all(less(out, array(1.0f))).item<bool>());
CHECK(all(greater_equal(out, array(0.0f))).item<bool>());