Error calculating gradient on function imported from R using reticulate

Question

I'm working on a problem right now where I am trying to use the optimizers from Tensorflow probability in Python to solve a simple optimization problem I've already defined in R.

Here are the steps:

Step 1: Define the original Python problem for solving the Rosenbrock banana function:

import contextlib
import functools
import os
import time

import numpy as np
import pandas as pd
import scipy as sp
from six.moves import urllib
from sklearn import preprocessing

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()

import tensorflow_probability as tfp


def make_val_and_grad_fn(value_fn):
    @functools.wraps(value_fn)
    def val_and_grad(x):
        return tfp.math.value_and_gradient(value_fn, x)
    return val_and_grad 

@contextlib.contextmanager
def timed_execution():
    t0 = time.time()
    yield
    dt = time.time() - t0
    print('Evaluation took: %f seconds' % dt)

def np_value(tensor):
  """Get numpy value out of possibly nested tuple of tensors."""
  if isinstance(tensor, tuple):
    return type(tensor)(*(np_value(t) for t in tensor))
  else:
    return tensor.numpy()

def run(optimizer):
    optimizer()  # Warmup.
    with timed_execution():
        result = optimizer()
    return np_value(result)
  
  
def run(optimizer):
    optimizer()  # Warmup.
    with timed_execution():
        result = optimizer()
    return np_value(result)
  

@make_val_and_grad_fn
def rosenbrock_test(coord):
    x, y = coord[..., 0], coord[..., 1]
    return((5.0-x)**2 + 10.0 * (y-x**2)**2)
  
optim_results = tfp.optimizer.lbfgs_minimize(
      rosenbrock_test,
      initial_position=rosenbrock_start, 
      tolerance=1e-12)
rosenbrock_start = tf.constant([2.,2.])



optim_results = tfp.optimizer.lbfgs_minimize(
      rosenbrock_test,
      initial_position=rosenbrock_start, 
      tolerance=1e-12)
rosenbrock_start = tf.constant([2.,2.])



print('L-BFGS Results')
print('Converged:', optim_results.converged)
print('Location of the minimum:', optim_results.position)
print('Number of iterations:', optim_results.num_iterations) 

Step 2: Define an identical function in R:

rosenbrock_for_r <- function(coord){
  x <- coord[1]
  y <- coord[2]
  return(   (5-x)^2 + 10 * (y-x^2)^2       )    }

rosenbrock_for_r(c(2,2))

Step 3: Define Python wrapper for the R function:

def rosenbrock_R(coord): 
    return(r.rosenbrock_for_r(coord))

The error occurs at this step:

temp = [2.0,2.0] 
tfp.math.value_and_gradient(rosenbrock_R, [2.,2.])

The error is:

TypeError: rosenbrock_R() takes 1 positional argument but 2 were given

I've tried investigating if I'm inputting something incorrectly to the function, but the implementation is the same as my native implementation:

def rosenbrock_alt(coord):
    x, y = coord[..., 0], coord[..., 1]
    return((5.0-x)**2 + 10.0 * (y-x**2)**2)  
  
temp = tf.constant([2.0,2.0])
  
tfp.math.value_and_gradient(rosenbrock_alt,temp)

This produces the expected output:

(<tf.Tensor: shape=(), dtype=float32, numpy=49.0>, <tf.Tensor: shape=(2,), dtype=float32, numpy=array([154., -40.], dtype=float32)>)

Answer 1

tfp.math.value_and_gradient will unpack the list into multiple arguments and diff with respect to each of them. You'll have to wrap in np.array or tf.convert_to_tensor .

Also, it's unclear how you will get a gradient for rosenbrock_for_r . You may have to use something like

@tf.custom_gradient
def f(x):
  def df(df_x):
    return r.grad_rosenbrock(x, df_x)
  return r.rosenbrock_for_r(x), df  # or x.numpy() but that will be eager-only

You could use tf.py_function to embed eager/r code into a TF graph.