| 11 | |
| 12 | |
| 13 | class TestJAX(unittest.TestCase): |
| 14 | def tanh(self, x): |
| 15 | y = np.exp(-2.0 * x) |
| 16 | return (1.0 - y) / (1.0 + y) |
| 17 | |
| 18 | def test_grad(self): |
| 19 | grad_tanh = grad(self.tanh) |
| 20 | ag = grad_tanh(1.0) |
| 21 | self.assertEqual(0.4199743, ag) |
| 22 | |
| 23 | def test_backend(self): |
| 24 | expected_backend = 'cpu' if not isGPU() else 'gpu' |
| 25 | self.assertEqual(expected_backend, jax.default_backend()) |
nothing calls this directly
no outgoing calls
no test coverage detected