Method _test_two_input_function

test/null/test_gradient.py:25–45

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

      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
    x = UOp.variable('x', -math.inf, math.inf, dtype=dtypes.float)
    y = UOp.variable('y', -math.inf, math.inf, dtype=dtypes.float)
    grads = compute_gradient(f(x, y), UOp.const(dtypes.float, 1.0), set([x, y]))
    gx, gy = grads[x], grads[y]

    for valx in [-5., -2.0, 0.0, 2.0, 5.]:
      for valy in [-5., -2.0, 0.0, 2.0, 5.]:
        # Substitute the values into the gradient expressions
        substitutions = {x: x.const_like(valx), y: y.const_like(valy)}
        tg_out_x = gx.substitute(substitutions).ssimplify()
        tg_out_y = gy.substitute(substitutions).ssimplify()

        tx = torch.tensor([valx], dtype=torch.float, requires_grad=True)
        ty = torch.tensor([valy], dtype=torch.float, requires_grad=True)
        torch_grad = torch.autograd.grad(jf(tx, ty), [tx, ty])
        torch_out_x, torch_out_y = [x.item() for x in torch_grad]

        self._cmp_nan_okay(tg_out_x, torch_out_x)
        self._cmp_nan_okay(tg_out_y, torch_out_y)

  # unary ops unit
  def test_recip(self): self._test_one_input_function(lambda x: 1.0/x)