tf_function.py

# First you might want to install & import tf 2.0
!pip install -q tensorflow==2.0.0-alpha0
import tensorflow as tf

# Let's implement a simple Linear layer (mx+b) using tf.function 
@tf.function
def add(a,b):
  return a + b

@tf.function
def linear_layer(m,x,b):
  return add(tf.matmul(m,x), b) # mx + b

linear_layer(tf.ones([2,2]), tf.ones([2,2]), tf.ones([2,2]))