Method _test_one_input_function

test/null/test_gradient.py:14–23

(self, f:Callable, jf:Callable|None=None)

    self.assertAlmostEqual(x, y, places=5)

  def _test_one_input_function(self, f:Callable, jf:Callable|None=None):
    if jf is None: jf = f
    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
    gx = compute_gradient(f(x), UOp.const(dtypes.float, 1.0), set([x]))[x]

    for val in [-5., -2.0, 0.0, 2.0, 5.]:
      tg_out = gx.substitute({x: x.const_like(val)}).ssimplify()
      tx = torch.tensor([val], dtype=torch.float, requires_grad=True)
      torch_out = torch.autograd.grad(jf(tx), tx)[0].item()
      self._cmp_nan_okay(tg_out, torch_out)

  def _test_two_input_function(self, f:Callable, jf:Callable|None=None):
    if jf is None: jf = f