random.h

Go to the documentation of this file.

// Copyright © 2023 Apple Inc.
 
#pragma once
 
#include <chrono>
#include <optional>
 
#include "mlx/array.h"
#include "mlx/stream.h"
#include "mlx/utils.h"
 
namespace mlx::core::random {
 
class KeySequence {
 public:
  explicit KeySequence(uint64_t seed);
 
  void seed(uint64_t seed);
  array next();
 
  // static default
  static KeySequence& default_() {
    static KeySequence ks(get_current_time_seed());
    return ks;
  }
 
 private:
  array key_;
  static uint64_t get_current_time_seed() {
    auto now = std::chrono::system_clock::now();
    return std::chrono::duration_cast<std::chrono::milliseconds>(
               now.time_since_epoch())
        .count();
  }
};
 
array key(uint64_t seed);
 
void seed(uint64_t seed);
 
array bits(
    const std::vector<int>& shape,
    int width,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
inline array bits(
    const std::vector<int>& shape,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return bits(shape, 4, key, s);
}
 
std::pair<array, array> split(const array& key, StreamOrDevice s = {});
 
array split(const array& key, int num, StreamOrDevice s = {});
 
array uniform(
    const array& low,
    const array& high,
    const std::vector<int>& shape,
    Dtype dtype = float32,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
template <typename T, typename U>
array uniform(
    T low,
    U high,
    const std::vector<int>& shape,
    Dtype dtype = float32,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return uniform(array(low), array(high), shape, dtype, key, to_stream(s));
}
 
array uniform(
    const std::vector<int>& shape,
    Dtype dtype,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
inline array uniform(
    const std::vector<int>& shape,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return uniform(shape, float32, key);
}
 
array normal(
    const std::vector<int>& shape,
    Dtype dtype,
    const float loc,
    const float scale,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
inline array normal(
    const std::vector<int>& shape,
    const float loc,
    const float scale,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return normal(shape, float32, loc, scale, key, s);
}
inline array normal(
    const std::vector<int>& shape,
    const Dtype dtype,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return normal(shape, dtype, 0.0, 1.0, key, s);
}
inline array normal(
    const std::vector<int>& shape,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return normal(shape, float32, 0.0, 1.0, key, s);
}
 
array multivariate_normal(
    const array& mean,
    const array& cov,
    const std::vector<int>& shape,
    Dtype dtype,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array randint(
    const array& low,
    const array& high,
    const std::vector<int>& shape,
    Dtype dtype = int32,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
template <typename T, typename U>
array randint(
    T low,
    U high,
    const std::vector<int>& shape,
    Dtype dtype = int32,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return randint(array(low), array(high), shape, dtype, key, to_stream(s));
}
 
array bernoulli(
    const array& p,
    const std::vector<int>& shape,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
array bernoulli(
    const array& p,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
template <typename T>
array bernoulli(
    T p,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return bernoulli(array(p), key, s);
}
 
template <typename T>
array bernoulli(
    T p,
    const std::vector<int>& shape,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return bernoulli(array(p), shape, key, s);
}
 
array bernoulli(
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array truncated_normal(
    const array& lower,
    const array& upper,
    const std::vector<int>& shape,
    Dtype dtype = float32,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array truncated_normal(
    const array& lower,
    const array& upper,
    Dtype dtype = float32,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array gumbel(
    const std::vector<int>& shape,
    Dtype dtype = float32,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array categorical(
    const array& logits,
    int axis,
    const std::vector<int>& shape,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array categorical(
    const array& logits_,
    int axis,
    int num_samples,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array categorical(
    const array& logits,
    int axis = -1,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
array laplace(
    const std::vector<int>& shape,
    Dtype dtype,
    const float loc,
    const float scale,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
inline array laplace(
    const std::vector<int>& shape,
    const float loc,
    const float scale,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return laplace(shape, float32, loc, scale, key, s);
}
inline array laplace(
    const std::vector<int>& shape,
    const Dtype dtype,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return laplace(shape, dtype, 0.0, 1.0, key, s);
}
inline array laplace(
    const std::vector<int>& shape,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {}) {
  return laplace(shape, float32, 0.0, 1.0, key, s);
}
 
/* Randomly permute the elements of x along the given axis. */
array permutation(
    const array& x,
    int axis = 0,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
/* A random permutation of `arange(x)` */
array permutation(
    int x,
    const std::optional<array>& key = std::nullopt,
    StreamOrDevice s = {});
 
} // namespace mlx::core::random