test_scipy_tf_function.py

import gym
import tensorflow as tf
from sklearn import preprocessing

from softlearning.utils.tensorflow import cast_and_concat


feature_fn = preprocessing.PolynomialFeatures(2)


def polynomial_inputs(x):
    return feature_fn.fit_transform(x)


@tf.function(experimental_relax_shapes=True)
def tf_function(inputs):
    inputs = cast_and_concat(inputs)
    y = tf.numpy_function(
        feature_fn.fit_transform,
        [inputs],
        tf.float32)
    return y


env = gym.envs.make('Hopper-v3')
observations = env.reset()[None, ...]
actions = env.action_space.sample()[None, ...]
inputs = {'observations': observations, 'actions': actions}
outputs = tf_function(inputs)