TypeError: An op outside of the function building code is being passed a "Graph" tensor. .. for Stax Tests

[DarrenZhang01]

As titled.

Traceback (most recent call last):
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py", line 3204, in equal
    _result = pywrap_tfe.TFE_Py_FastPathExecute(
tensorflow.python.eager.core._FallbackException: This function does not handle the case of the path where all inputs are not already EagerTensors.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "stax_test.py", line 1074, in test_leaky_relu
    self.assertAllClose(kernels_leaky_relu, kernels_ab_relu)
  File "/Users/Darren/Desktop/SchoolYear/2020Summer/GSoC_TensorFlow/GitHub/TensorFlow_GSoC/neural-tangents/neural_tangents/utils/test_utils.py", line 152, in assertAllClose
    self.assertEqual(x_dict[field.name], y_dict[field.name])
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/unittest/case.py", line 912, in assertEqual
    assertion_func(first, second, msg=msg)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/unittest/case.py", line 902, in _baseAssertEqual
    if not first == second:
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py", line 1491, in tensor_equals
    return gen_math_ops.equal(self, other, incompatible_shape_error=False)
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py", line 3211, in equal
    return equal_eager_fallback(
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py", line 3247, in equal_eager_fallback
    _result = _execute.execute(b"Equal", 1, inputs=_inputs_flat, attrs=_attrs,
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/eager/execute.py", line 75, in quick_execute
    raise e
  File "/Library/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/tensorflow/python/eager/execute.py", line 59, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: Identity:0