Thanks for your answer.
I have tried to use jit to speed up the computation. It doesn’t work. I have been able to jit the original function, but this does not lead to a relevant speedup. I can’t jit the gradients and the Hessian:

  
      gist.github.com
  
  
    https://gist.github.com/darteaga/8df73f503518de9c8c783d11bae4e6ee 6
compare_ipyopt_jax.ipynb
{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
This file has been truncated. show original




  
  
    
    
  
  


In the same gist I compare also to JAX 7 for gradient and hessian. With JAX it is possible to precompute the gradients using jit. It seems that PyTorch is more or less at the level as JAX for the gradient computation, but not for Hessian computation: Hessian is about 20x slower in PyTorch than in JAX.