45 lines
1.1 KiB
Go
45 lines
1.1 KiB
Go
|
//go:build darwin && arm64 && mlx
|
||
|
|
|
||
|
|
package mlx
|
||
|
|
|
||
|
|
/*
|
||
|
|
#include "mlx/c/mlx.h"
|
||
|
|
*/
|
||
|
|
import "C"
|
||
|
|
|
||
|
|
// RandomCategorical samples from a categorical distribution defined by logprobs.
|
||
|
|
// Returns indices sampled according to the log-probability distribution along the last axis.
|
||
|
|
func RandomCategorical(logprobs *Array) *Array {
|
||
|
|
out := New("RANDOM_CATEGORICAL", logprobs)
|
||
|
|
// shape for output: same as input but last dim removed
|
||
|
|
C.mlx_random_categorical_shape(
|
||
|
|
&out.ctx,
|
||
|
|
logprobs.ctx,
|
||
|
|
C.int(-1), // axis
|
||
|
|
nil, C.int(0), // empty shape = infer from input
|
||
|
|
nil, // key (use default)
|
||
|
|
DefaultStream().ctx,
|
||
|
|
)
|
||
|
|
return out
|
||
|
|
}
|
||
|
|
|
||
|
|
// RandomUniform generates uniform random values in [low, high).
|
||
|
|
func RandomUniform(low, high float32, shape []int32, dtype DType) *Array {
|
||
|
|
out := New("RANDOM_UNIFORM")
|
||
|
|
cShape := make([]C.int, len(shape))
|
||
|
|
for i, s := range shape {
|
||
|
|
cShape[i] = C.int(s)
|
||
|
|
}
|
||
|
|
lo := FromValue(low)
|
||
|
|
hi := FromValue(high)
|
||
|
|
C.mlx_random_uniform(
|
||
|
|
&out.ctx,
|
||
|
|
lo.ctx, hi.ctx,
|
||
|
|
&cShape[0], C.int(len(cShape)),
|
||
|
|
C.mlx_dtype(dtype),
|
||
|
|
nil, // key
|
||
|
|
DefaultStream().ctx,
|
||
|
|
)
|
||
|
|
return out
|
||
|
|
}
|