def testScipySpecialFun(self, scipy_op, lax_op, rng_factory, shapes,
dtypes, test_autodiff, nondiff_argnums):
rng = rng_factory(self.rng())
args_maker = self._GetArgsMaker(rng, shapes, dtypes)
args = args_maker()
self.assertAllClose(scipy_op(*args),
lax_op(*args),
atol=1e-3,
rtol=1e-3,
check_dtypes=False)
self._CompileAndCheck(lax_op, args_maker, rtol=1e-4)
if test_autodiff:
def partial_lax_op(*vals):
list_args = list(vals)
for i in nondiff_argnums:
list_args.insert(i, args[i])
return lax_op(*list_args)
assert list(nondiff_argnums) == sorted(set(nondiff_argnums))
diff_args = [
x for i, x in enumerate(args) if i not in nondiff_argnums
]
jtu.check_grads(partial_lax_op,
diff_args,
order=1,
atol=jtu.if_device_under_test("tpu", .1, 1e-3),
rtol=.1,
eps=1e-3)
def testScipySpecialFun(self, scipy_op, lax_op, rng_factory, shapes,
dtypes, test_autodiff, nondiff_argnums):
if (jtu.device_under_test() == "cpu"
and (lax_op is lsp_special.gammainc
or lax_op is lsp_special.gammaincc)):
# TODO(b/173608403): re-enable test when LLVM bug is fixed.
raise unittest.SkipTest("Skipping test due to LLVM lowering bug")
rng = rng_factory(self.rng())
args_maker = self._GetArgsMaker(rng, shapes, dtypes)
args = args_maker()
self.assertAllClose(scipy_op(*args),
lax_op(*args),
atol=1e-3,
rtol=1e-3,
check_dtypes=False)
self._CompileAndCheck(lax_op, args_maker, rtol=1e-4)
if test_autodiff:
def partial_lax_op(*vals):
list_args = list(vals)
for i in nondiff_argnums:
list_args.insert(i, args[i])
return lax_op(*list_args)
assert list(nondiff_argnums) == sorted(set(nondiff_argnums))
diff_args = [
x for i, x in enumerate(args) if i not in nondiff_argnums
]
jtu.check_grads(partial_lax_op,
diff_args,
order=1,
atol=jtu.if_device_under_test("tpu", .1, 1e-3),
rtol=.1,
eps=1e-3)
def testScipySpecialFun(self, scipy_op, lax_op, rng_factory, shapes,
dtypes, test_autodiff):
rng = rng_factory(self.rng())
args_maker = self._GetArgsMaker(rng, shapes, dtypes)
args = args_maker()
self.assertAllClose(scipy_op(*args),
lax_op(*args),
atol=1e-3,
rtol=1e-3,
check_dtypes=False)
self._CompileAndCheck(lax_op, args_maker, check_dtypes=True, rtol=1e-5)
if test_autodiff:
jtu.check_grads(lax_op,
args,
order=1,
atol=jtu.if_device_under_test("tpu", .1, 1e-3),
rtol=jtu.if_device_under_test("tpu", .1, .02),
eps=1e-3)
