Issue vectorising mlx.core.random.split using mlx.core.vmap

[fbadine]

Hi,

I am trying to use mlx.core.vmap in order to vectorise an implementation of the gamma distribution and part of that I want to vectorise splitting the PRNG keys.

I have the following code simplified to show the error I am getting:

import math
import numpy as np
import mlx.core as mx

def _split(key, split_count_array):
    return mx.random.split(key, split_count_array.item())

a = 5.
shape = (10,)
seed = 1337

key = mx.random.key(1337)
print(key)

a = mx.broadcast_to(a, shape)
a_shape = np.shape(a)

split_count = math.prod(a_shape[key.ndim:])
split_count_array = mx.array(split_count)
keys = mx.vmap(_split, in_axes=(0, None))(key, split_count_array)

The error I am receiving when I run this code is the following:

Traceback (most recent call last):
File "/Users/fadibadine/Documents/Projects/mlx/test_vmap.py", line 20, in
keys = mx.vmap(_split, in_axes=(0, None))(key, split_count_array)
File "/Users/fadibadine/Documents/Projects/mlx/test_vmap.py", line 6, in _split
return mx.random.split(key, split_count_array.item())
ValueError: Expected key shape (2) but received ().

I do not know what the problem is and why key doesn't seem to be passed to function _split

Thanks!

[awni]

I can't quite tell what you are trying to accomplish from that snippet of code. But one red flag is the fact that you are taking the item of your split_count_array which is going to eval the graph and get you a scalar. You should not do that inside a function transformation like VMAP (it's pretty much always a bug to do so).

You can split the key safely inside the vmap function but you should hardcode the number of splits. The vmap over that will take care of producing the number of RNG keys you need for each example in the batch.

Maybe you can share a little more of what you are trying to do so we can help more?

[fbadine]

Thanks for your reply. In fact the error is about key not being passed it seems.
Anyway, what I am trying to achieve is gamma distribution for Keras when the backend is set to MLX. MLX, so far, does not have a gamma distribution function that Keras API can call so I am trying to implement one in the Keras backend.
Here is the code. To simplify, I used mlx functions only (not the higher level ones created in Keras MLX. There are though 2 functions that are part of Keras MLX branch while_loop and convert_to_tensor

Here is the full code:

import numpy as np
import math
import mlx.core as mx

def _gamma_one(alpha, key):
    def _const(example, val):
        dtype = example.dtype
        return convert_to_tensor(val).astype(dtype)
    
    zero = _const(alpha, 0)
    one = _const(alpha, 1)
    minus_one = _const(alpha, -1)
    one_over_two = _const(alpha, 0.5)
    one_over_three = _const(alpha, 1./3)
    squeeze_const = _const(alpha, 0.0331)
    dtype = alpha.dtype


    boost_mask = mx.greater_equal(alpha, one)
    alpha_orig = alpha
    alpha = mx.where(boost_mask, alpha, mx.add(alpha, one))

    d = mx.subtract(alpha, one_over_three)
    c = mx.divide(one_over_three, mx.sqrt(d))

    def _cond_fn(k, X, V, U):
        condition = mx.logical_and(
            mx.greater_equal(
                U, 
                mx.subtract(
                    one,
                    mx.multiply(
                        squeeze_const,
                        mx.multiply(X, X)
                    )
                )
            ),
            mx.greater_equal(
                mx.log(U),
                mx.add(
                    mx.multiply(X, one_over_two),
                    mx.multiply(
                        d, 
                        mx.add(
                            mx.subtract(one, V),
                            mx.log(V)
                        )
                    )
                )
            )
        )

        return condition

    def _body_fn(k, X, V, U):
        def _next_kxv(k, x, v):
            key = k

            key, subkey = mx.random.split(key)
            x = mx.random.normal(shape=(), dtype=dtype, key=subkey)
            v = mx.add(one, mx.multiply(x, c))
            return key, x, v
        
        def _lambda_fn(k, x, v):
            cond = mx.less_equal(v, zero)

            return cond
        
        key = k
        key, x_key, U_key = mx.random.split(key, 3)
        _, x, v = while_loop(
            lambda k, x, v: mx.less_equal(v, zero), _next_kxv, (x_key, zero, minus_one)
            #_lambda_fn, _next_kxv, (x_key, zero, minus_one)
        )
        X = mx.multiply(x, x)
        V = mx.multiply(mx.multiply(v, v), v)
        U = mx.random.uniform(low=0.0, high=1.0, shape=(), dtype=dtype, key=U_key)
        return key, X, V, U
    
    key, subkey = mx.random.split(key)
    _, _, V, _ = while_loop(_cond_fn, _body_fn, (key, zero, one, _const(alpha, 2)))

    samples = 1 - mx.random.uniform(low=0.0, high=1.0, shape=(), dtype=dtype, key=subkey)
    boost = mx.where(boost_mask, one, mx.power(samples, mx.divide(one, alpha_orig)))

    return mx.multiply(mx.multiply(d, V), boost)

def _gamma_impl(a, key):
    a_shape = np.shape(a)
    split_count = math.prod(a_shape[key.ndim:])
    keys = key.flatten()

    keys = mx.vmap(mx.random.split, in_axes=(0, None))(keys, split_count)
    keys = keys.flatten()

    alphas = a.flatten()

    samples = mx.vmap(_gamma_one)(alphas, keys)

    return samples

def _gamma(a, key, shape, dtype):
    if shape is None:
        shape = np.shape(a)

    if np.shape(a) != shape:
        a = mx.broadcast_to(a, shape)
    a = a.astype(dtype) 

    if np.shape(a) != shape:
        a = mx.broadcast_to(a, shape)

    return _gamma_impl(a, key)

def gamma(shape, alpha, dtype=None, seed=None):
    dtype = dtype or "float32"

    if not mx.issubdtype(dtype, mx.floating):
        raise ValueError(f"dtype argument to `gamma` must be a float "
                         f"dtype, got {dtype}")

    key = mx.random.key(seed)

    return _gamma(alpha, key=key, shape=shape, dtype=dtype)

[awni]

Cool, thanks for sharing the code! There is a bit more than I was expecting 😅 .

As for the key, if you want to simplify it you could remove the key entirely and just seed the global PRNG. I'm not sure if that works with the semantics of the seed kwarg in the gamma function.

If that doesn't work, here is a very simple example of how to split the keys with a vmap:

# Generate (5, 5) array of uniform numbers 
def uniform(key):
    return mx.random.uniform(shape=(5, 5), key=key)

key = mx.random.key(0)

# shape (5, 5)
samples = uniform(key)
print(samples.shape)

key = mx.random.split(key, 100)
# Shape (100, 5 , 5)
samples = mx.vmap(uniform)(key)
print(samples.shape)