I’m optimizing the dual of a convex optimization problem, that yields an unconstrained optimization problem of the above form. My dual variable is of size n_data x m since originally I had n_data x m constraints in my primal constrained convex problem and each row of my dual variable corresponds to a specific datapoint. I would prefer batching these variables with the datapoints to allow for easy shuffling.