Remove @tf.function wrappers

Author: Squadrick

tensorflow_addons/text/crf.py

@@ -20,8 +20,9 @@
 import numpy as np
 import tensorflow as tf
 
+# TODO: Wrap functions in @tf.function once
+# https://github.com/tensorflow/tensorflow/issues/29075 is resolved
 
-@tf.function
 def crf_sequence_score(inputs, tag_indices, sequence_lengths,
                        transition_params):
     """Computes the unnormalized score for a tag sequence.
@@ -66,7 +67,6 @@ def _multi_seq_fn():
         return _multi_seq_fn()
 
 
-@tf.function
 def crf_multitag_sequence_score(inputs, tag_bitmap, sequence_lengths,
                                 transition_params):
     """Computes the unnormalized score of all tag sequences matching
@@ -115,7 +115,6 @@ def _multi_seq_fn():
         return _multi_seq_fn()
 
 
-@tf.function
 def crf_log_norm(inputs, sequence_lengths, transition_params):
     """Computes the normalization for a CRF.
 
@@ -163,7 +162,6 @@ def _multi_seq_fn():
         return _multi_seq_fn()
 
 
-@tf.function
 def crf_log_likelihood(inputs,
                        tag_indices,
                        sequence_lengths,
@@ -201,7 +199,6 @@ def crf_log_likelihood(inputs,
     return log_likelihood, transition_params
 
 
-@tf.function
 def crf_unary_score(tag_indices, sequence_lengths, inputs):
     """Computes the unary scores of tag sequences.
 
@@ -236,7 +233,6 @@ def crf_unary_score(tag_indices, sequence_lengths, inputs):
     return unary_scores
 
 
-@tf.function
 def crf_binary_score(tag_indices, sequence_lengths, transition_params):
     """Computes the binary scores of tag sequences.
 
@@ -272,7 +268,6 @@ def crf_binary_score(tag_indices, sequence_lengths, transition_params):
     return binary_scores
 
 
-@tf.function
 def crf_forward(inputs, state, transition_params, sequence_lengths):
     """Computes the alpha values in a linear-chain CRF.
 
@@ -364,7 +359,6 @@ def call(self, inputs, state):
         return backpointers, new_state
 
 
-@tf.function
 def crf_decode_forward(inputs, state, transition_params, sequence_lengths):
     mask = tf.sequence_mask(sequence_lengths, tf.shape(inputs)[1])
     crf_fwd_cell = CrfDecodeForwardRnnCell(transition_params)
@@ -373,7 +367,6 @@ def crf_decode_forward(inputs, state, transition_params, sequence_lengths):
     return crf_fwd_layer(inputs, state, mask=mask)
 
 
-@tf.function
 def crf_decode_backward(inputs, state):
     """Computes backward decoding in a linear-chain CRF."""
     inputs = tf.transpose(inputs, [1, 0, 2])
@@ -387,7 +380,6 @@ def _scan_fn(state, inputs):
     return tf.transpose(tf.scan(_scan_fn, inputs, state), [1, 0, 2])
 
 
-@tf.function
 def crf_decode(potentials, transition_params, sequence_length):
     """Decode the highest scoring sequence of tags in TensorFlow.