Convex Hull

[pharringtonp19]

I would be grateful for any feedback on how to improve the following bit of code.

I am looking to take a matrix, $X$, randomly sample a row, $x_i$, and determine if the randomly sampled vector is in the convex hull formed by the remaining rows: $x_i \in Cov(X_{-i})$

The parameter $d$ allows me to restrict the number of columns I consider.

@partial(jax.jit, static_argnums=(2, 3))
def split(key, X, n, d):
    idx = jax.random.permutation(key, jnp.arange(X.shape[0]))
    X = X[idx]
    return X[:n,:d], X[n,:d]

def check_convex_hull(X, p):
    b = jnp.hstack((jnp.ones(shape=(1,)),p))
    X = jax.device_put(X, sharding)
    Q = jnp.zeros(shape=(X.shape[0], X.shape[0]))
    c = jnp.zeros(shape=(X.shape[0],))
    A = jnp.vstack((jnp.ones(shape=(X.shape[0],)), X.T))
    G = -jnp.eye(X.shape[0])
    h = jnp.zeros(shape=(X.shape[0],))
    qp = OSQP()
    sol = qp.run(
                init_params=None,
                params_obj=(Q, c), 
                params_eq=(A, b), 
                params_ineq=(G, h)).params.primal
    
    return jnp.isclose(A @sol, b, rtol=1e-3) 
    

@partial(jax.jit, static_argnums=(2, 3))
def Run(key, X, n, d):
    X, p = split(key, X, n, d)
    return check_convex_hull(X, p)