def grads(x, y):
return (
grad(f, 'x')(x, y),
grad(f, 0)(x, y),
grad(f, 'y')(x, y),
grad(f, 'x', 'y')(x, y),
grad(f, '*')(x, y),
value_and_grad(f)(x, y),
)
def test_grad_reduce(xs, ys):
def f(xs, ys):
return array_reduce(scalar_add, xs * ys, ())
return grad(f)(xs, ys)
def test_grad_cast(x):
def f(x):
return scalar_cast(x, f16)
return grad(f)(x)
def test_grad(x, y):
def f(x, y):
return x * (y + x)
return grad(f)(x, y)
def step(model, x, y):
return grad(cost)(model, x, y)
def test_backward_infer(model, x, y):
return grad(cost)(model, x, y)
def gradbad3(x, y, z):
return grad(f, z)(x, y)
def gradbad2(x, y):
return grad(f, 'z')(x, y)
def gradbad(x, y):
return grad(f, (0, 1))(x, y)
