How to compute in tensorflow the partial derivative of each observation from the output w.r.t each input observation

Question

I'm using a sequential FFNN to price financial options, and I'm trying to compute the option deltas. While I've achieved better accuracy in option pricing (the output of the model) compared to the Black-Scholes formula, my computed deltas are less accurate. Can anyone help me understand why my neural network-based model produces more accurate option prices but less accurate deltas than the Black-Scholes formula? I know I should use tape.jacobian function but it keeps killing my kernel. Instead of I tried to use the tape.gradient function but the deltas are less accurate than BS formula. I've read from tensorflow that it computes the partial derivative as the sum of ys. What I need is the partial derivative of each output observation w.r.t to each input observation. I use: Python 3.11.0 Tensorflow 2.12.04 Numpy 1.23.5 I've already found some posts but none of these answer to my question.

Here is my code (X_test is a constant containing 4 features including the option price)

df_op = df_op[['C0/K', 'M', 'Time_to_expiry', 'garch_vol', 'cp_int']]

def split_X_Y(data):
    X = data.copy()
    y = X.pop('C0/K')
    return X, y


def split_data(X, y):
    X_train_full, X_test, y_train_full, y_test = train_test_split(X, y, test_size=0.20, random_state=42)
    X_train, X_valid, y_train, y_valid = train_test_split(X_train_full, y_train_full, test_size=0.25, random_state= 42)
    X_index = X_test.index.to_frame(index=True)
    return X_test, y_test, X_train, y_train, X_valid, y_valid, X_index


    

#Split X and y
X, y = split_X_Y(data=df_op)

#Cross_validation
X_test, y_test, X_train, y_train, X_valid, y_valid, X_index = split_data(X, y)


X_test = tf.constant(X_test, dtype=tf.float32)
y_test = tf.constant(y_test, dtype=tf.float32)
X_train = tf.constant(X_train, dtype=tf.float32)
y_train = tf.constant(y_train, dtype=tf.float32)
X_valid = tf.constant(X_valid, dtype=tf.float32)
y_valid = tf.constant(y_valid, dtype=tf.float32)

model = build_ANN(X_test, num_units = 30, dropout_rate = 0.25, loss= 'mse')


history, learning_curves = history(model, X_train, X_valid, y_train, y_valid)

def compute_jacobian(X, model, X_index):
    soft_memory_limit_mb = 800  

    try:
        with tf.GradientTape(persistent=True) as tape:

            tape.watch(X)

            y_pred = model(X)

        # Check if memory usage exceeds the threshold
        if is_memory_limit_exceeded(soft_memory_limit_mb):
            raise MemoryError("Memory usage exceeded soft limit during tape.watch")

        # Computes the gradients of predictions with respect to input features
        jacobian = tape.jacobian(y_pred, X)
        del tape
        jacobian_reshaped = jacobian.numpy().reshape(X.shape[0], -1)
        column_names = [f'Jacobian_{i}' for i in range(jacobian_reshaped.shape[1])]
        jacobian_df = pd.DataFrame(jacobian_reshaped, columns=column_names)

        def shift_non_zero(row):
            non_zero_values = [value for value in row if value != 0]
            return non_zero_values + [0] * (len(row) - len(non_zero_values))
        
        jacobian_df = jacobian_df.apply(shift_non_zero, axis=1, result_type='expand')
        jacobian_df = jacobian_df.loc[:, (jacobian_df != 0).any()]
        jacobian_df = jacobian_df.iloc[:, [0]]
        jacobian_df.rename(columns={jacobian_df.columns[0]: 'delta_nn'}, inplace=True)
        jacobian_df.index = X_index.index

    except MemoryError:
        jacobian_df = compute_jacobian_alternative(X, model, X_index)

    return jacobian_df


def compute_jacobian_alternative(X, model, X_index):
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(X)
        y_pred = model(X)
    gradients2 = tape.gradient(y_pred, X)
    del tape

    gradient_np = gradients2.numpy()
    jacobian_df = pd.DataFrame(gradient_np)
    jacobian_df = jacobian_df.iloc[:, [0]]
    jacobian_df.rename(columns={jacobian_df.columns[0]: 'delta_nn'}, inplace=True)
    jacobian_df.index = X_index.index

    return jacobian_df