tfl.test_utils.assert_sonnet_equivalent_to_keras

Asserts that a Sonnet module is "equivalent" to a Keras layer.

tfl.test_utils.assert_sonnet_equivalent_to_keras(
    test,
    sonnet_module_ctor,
    keras_layer_ctor,
    training_inputs,
    training_labels,
    epsilon=0.0001
)

Creates a Sonnet module and a Keras layer using the given constructors. It then uses both models to evaluate the given 'training_inputs' tensor and asserts that the results are equal. It then trains both models and asserts that the final loss (w.r.t the given 'training_labels') and the post-training predictions of both models on 'training_inputs' are also equal.

Args
test	a tf.test.TestCase whose 'assert...' methods to use for assertion.
sonnet_module_ctor	A callable that takes no arguments and returns the Sonnet module to use.
keras_layer_ctor	A callable that takes no arguments and returns the Keras layer to use.
training_inputs	Tensor of shape (batch_size, ....) tensor containing the training inputs.
training_labels	Tensor of shape (batch_size, ....). tensor containing the training labels.
epsilon	float. Sensitivity of comparison. Comparison of model predictions and losses are done using test.assertNear and test.assertNDArrayNear. This is the value to pass as the 'err' parameter to these assertion methods.