def test_ensure_true(self):
    """Tests that we can check arbitrary functions in `EnsureTrue`."""

    def _nonzero(a):
      return tf.not_equal(a, 0)

    # TODO(jhclark): Notice that the syntax is currently insufficient for
    # handling tuple return types with `EnsureTrue`.
    @tc.contract(tc.EnsureTrue(_nonzero, tensors=["RESULT"], error="Fail."))
    def my_func(a: tf.Tensor) -> tf.Tensor:
      return a

    a = tf.constant([1, 2])
    result = my_func(a)
    with self.session() as sess:
      # Passes:
      sess.run(result)

    a = tf.constant([0, 1])
    with self.assertRaises(tf.errors.InvalidArgumentError):
      result = my_func(a)
      with self.session() as sess:
        # Fails:
        sess.run(result)
Tests that we can check arbitrary functions in `EnsureTrue`.
language/canine/tensor_contracts_test.py