Hi,
thank you very much for your kind offer! I will definitely check the Jacobian function that you proposed, but in the meanwhile, I attached the code that returns my issue below:
import numpy as np
import torch


# --------------------------------------------------------------------------- #
# Parameters
# --------------------------------------------------------------------------- #
k = 0.1
a = 1

alpha = 0.36
beta = 0.95
psi = 0.25
theta = 1.5
s = 0.01
mu = 0
rho = 0.95

# Gauss-Hermite quadrature
# Nodes
x5 = np.sqrt(2) * s * np.array(
    [2.020182870456086, 0.9585724646138185, 0, -0.9585724646138185,
     -2.020182870456086]) + mu
# Weights
omega5 = np.pi**(-1/2) * torch.tensor(
    [0.01995324205904591, 0.3936193231522412, 0.9453087204829419,
     0.3936193231522412, 0.01995324205904591])

nvar = 1  # Number of variables
ncon = nvar  # Number of constraints

nnzj = int(nvar * ncon)  # Number of (possibly) non-zeros in Jacobian


def ls_compute(k, A, alpha=alpha, psi=psi, theta=theta):
    """ Return the optimal labor supply """
    return (((1-alpha) * A * k**alpha) / (psi*theta))**(1 / (alpha+theta-1))


# Current labor supply
ls = ls_compute(k=k, A=a, alpha=alpha, psi=psi, theta=theta)

# AR(1) shock
aplus = torch.empty(x5.shape)
for epsilon_idx, epsilon_plus in enumerate(x5):
    aplus[epsilon_idx] = a**rho * np.exp(epsilon_plus)


# --------------------------------------------------------------------------- #
# Functions
# eval_g returns the set of constraints that needs to be first-derivative
# eval_grad_f returns (1) the structure of the Jacobian and (2) its values
# --------------------------------------------------------------------------- #

def eval_g(x):
    """ The system of non-linear equilibrium conditions
    x[0]: Capital stock in the next period
    """
    assert len(x) == nvar

    # Consumption today
    con = a * k**alpha * ls**(1-alpha) - x[0]

    # Labor supply tomorrow
    # ls_plus = np.empty(x5.shape)
    ls_plus = torch.empty(x5.shape)
    for aplus_idx, aplus_val in enumerate(aplus):
        ls_plus[aplus_idx] = ls_compute(
            k=x[0], A=aplus_val, alpha=alpha, psi=psi, theta=theta)
    # Capital stock tomorrow
    k_plusplus = torch.empty(aplus.shape)
    for aplus_idx, aplus_val in enumerate(aplus):
        # state_plus = torch.tensor([x[0], aplus_val])[None, :].float()
        k_plusplus[aplus_idx] = torch.tensor(k)  # Needs to be modified

    # Consumption tomorrow
    con_plus = aplus * x[0]**alpha * ls_plus**(1-alpha) - k_plusplus

    # ----------------------------------------------------------------------- #
    # Euler equation
    # ----------------------------------------------------------------------- #
    g0 = 1 / con - beta * alpha * torch.sum(omega5 * (
        1 / con_plus * aplus * x[0]**(alpha-1) * ls_plus**(1-alpha)))

    return np.array([g0])


# Test
print(r"Constraint returns {}".format(eval_g(np.array([k]))))
# sys.exit(0)


def eval_jac_g(x, flag):
    """ Numerical approximation of the Jacobian of the system of
    non-linear equilibrium conditions
    Jacobian is computed by the automatic gradient of PyTorch """
    assert len(x) == nvar

    row_idx = np.empty(nnzj, dtype=int)  # Row index
    col_idx = np.empty(nnzj, dtype=int)  # Column index

    # Jacobian matrix structure
    if flag:
        for i in range(ncon):
            for j in range(nvar):
                row_idx[j + i * nvar] = i
                col_idx[j + i * nvar] = j

        return (row_idx, col_idx)

    else:
        # Automatic gradient by PyTorch
        assert len(x) == nvar
        _eval_jac_g = np.empty(nnzj)
        for i in range(ncon):
            x_tensor = torch.tensor(x, requires_grad=True)
            eval_g_tensor = eval_g(x_tensor)[i]
            grads = torch.autograd.grad(eval_g_tensor, x_tensor,
                                        create_graph=True)
            for j in range(nvar):
                _eval_jac_g[j + i*nvar] = grads[0][j].detach().numpy()
        return _eval_jac_g


# Test
print(r"The structure of the Jacobian is {}".format(
    eval_jac_g(np.array([k]), True)))
print(r"The values of each element in the Jacobian are {}".format(
    eval_jac_g(np.array([k]), False)))

The code should be standalone itself. Please note that in this script, the structure of Jacobian is flattened, i.e., it should be (df0/dx0, df0/dx1, …, df0/dxn, df1/dx0, …, dfn/dxn). In eval_jac_g, the first block returns the structure of Jacobian, and the second one returns its values. Basically you need to look at the second block…
I highly appreciate if you provide me the correct syntax to return Jacobian.
Thank you very much in advance.