vector<float> theta_data(num_elems);
  memcpy(theta_data.data(), TF_TensorData(theta_tensor),
         TF_TensorByteSize(theta_tensor));

  // Initialize space for the numerical gradient.
  vector<float> dtheta_approx(num_elems);

  // Get theta shape and store in theta_dims.
  int num_dims = TF_NumDims(theta_tensor);
  vector<int64_t> theta_dims(num_dims);