Function mlx_custom_function 

pub unsafe extern "C" fn mlx_custom_function(
    res: *mut mlx_closure,
    fun: mlx_closure,
    fun_vjp: mlx_closure_custom,
    fun_jvp: mlx_closure_custom_jvp,
    fun_vmap: mlx_closure_custom_vmap,
) -> c_int