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

Faster bits and bernoulli (#1535)

Browse Source 
* faster bits and bernoulli

* fix bernoulli
This commit is contained in:
Awni Hannun
2024-10-28 11:11:00 -07:00
committed by GitHub
parent 91f6c499d7
commit d3cd26820e
3 changed files with 14 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
mlx/backend/metal/kernels/random.metal
View File

@@ -34,8 +34,8 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
[[kernel]] void rbitsc(
device const uint32_t* keys,
device char* out,
device const bool& odd,
device const uint& bytes_per_key,
constant const bool& odd,
constant const uint& bytes_per_key,
uint2 grid_dim [[threads_per_grid]],
uint2 index [[thread_position_in_grid]]) {
auto kidx = 2 * index.x;
@@ -67,8 +67,8 @@ rbits threefry2x32_hash(const thread uint2& key, uint2 count) {
[[kernel]] void rbits(
device const uint32_t* keys,
device char* out,
device const bool& odd,
device const uint& bytes_per_key,
constant const bool& odd,
constant const uint& bytes_per_key,
constant const int& ndim,
constant const int* key_shape,
constant const size_t* key_strides,