mirror of
https://github.com/ROCm/jax.git
synced 2025-04-17 12:26:07 +00:00
ef5e4a4035
Instead, we skip tests that the PJRT C API doesn't support. We had this tag for feature development so it was easy to broadly disable, but now we don't expect to need to do that.
453 lines
13 KiB
Python
453 lines
13 KiB
Python
# Copyright 2022 The JAX Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import io
|
|
import re
|
|
import textwrap
|
|
import unittest
|
|
|
|
from typing import IO, Sequence, Tuple
|
|
|
|
from absl.testing import absltest
|
|
import jax
|
|
from jax.config import config
|
|
from jax.experimental import pjit
|
|
from jax._src import debugger
|
|
from jax._src import test_util as jtu
|
|
from jax._src import xla_bridge
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
|
|
config.parse_flags_with_absl()
|
|
|
|
def make_fake_stdin_stdout(commands: Sequence[str]) -> Tuple[IO[str], io.StringIO]:
|
|
fake_stdin = io.StringIO()
|
|
fake_stdin.truncate(0)
|
|
for command in commands:
|
|
fake_stdin.write(command + "\n")
|
|
fake_stdin.seek(0)
|
|
return fake_stdin, io.StringIO()
|
|
|
|
def _format_multiline(text):
|
|
return textwrap.dedent(text).lstrip()
|
|
|
|
prev_xla_flags = None
|
|
|
|
def setUpModule():
|
|
global prev_xla_flags
|
|
# This will control the CPU devices. On TPU we always have 2 devices
|
|
prev_xla_flags = jtu.set_host_platform_device_count(2)
|
|
|
|
# Reset to previous configuration in case other test modules will be run.
|
|
def tearDownModule():
|
|
prev_xla_flags()
|
|
|
|
foo = 2
|
|
|
|
class CliDebuggerTest(jtu.JaxTestCase):
|
|
|
|
def test_debugger_eof(self):
|
|
stdin, stdout = make_fake_stdin_stdout([])
|
|
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
with self.assertRaises(SystemExit):
|
|
f(2.)
|
|
jax.effects_barrier()
|
|
|
|
def test_debugger_can_continue(self):
|
|
stdin, stdout = make_fake_stdin_stdout(["c"])
|
|
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
f(2.)
|
|
jax.effects_barrier()
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
(jdb) """)
|
|
self.assertEqual(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_can_print_value(self):
|
|
stdin, stdout = make_fake_stdin_stdout(["p x", "c"])
|
|
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
(jdb) Array(2., dtype=float32)
|
|
(jdb) """)
|
|
f(jnp.array(2., jnp.float32))
|
|
jax.effects_barrier()
|
|
self.assertEqual(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_can_print_value_in_jit(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p x", "c"])
|
|
|
|
@jax.jit
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
(jdb) array(2., dtype=float32)
|
|
(jdb) """)
|
|
f(jnp.array(2., jnp.float32))
|
|
jax.effects_barrier()
|
|
self.assertEqual(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_can_print_multiple_values(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p x, y", "c"])
|
|
|
|
@jax.jit
|
|
def f(x):
|
|
y = x + 1.
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
(jdb) (array(2., dtype=float32), array(3., dtype=float32))
|
|
(jdb) """)
|
|
f(jnp.array(2., jnp.float32))
|
|
jax.effects_barrier()
|
|
self.assertEqual(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_can_print_context(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["l", "c"])
|
|
|
|
@jax.jit
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
f(2.)
|
|
jax.effects_barrier()
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) > .*debugger_test\.py\([0-9]+\)
|
|
@jax\.jit
|
|
def f\(x\):
|
|
y = jnp\.sin\(x\)
|
|
-> debugger\.breakpoint\(stdin=stdin, stdout=stdout, backend="cli"\)
|
|
return y
|
|
.*
|
|
\(jdb\) """)
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_can_print_backtrace(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["bt", "c"])
|
|
|
|
@jax.jit
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:.*
|
|
\(jdb\) Traceback:.*
|
|
""")
|
|
f(2.)
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_can_work_with_multiple_stack_frames(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["l", "u", "p x", "d", "c"])
|
|
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
|
|
@jax.jit
|
|
def g(x):
|
|
y = f(x)
|
|
return jnp.exp(y)
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) > .*debugger_test\.py\([0-9]+\)
|
|
def f\(x\):
|
|
y = jnp\.sin\(x\)
|
|
-> debugger\.breakpoint\(stdin=stdin, stdout=stdout, backend="cli"\)
|
|
return y
|
|
.*
|
|
\(jdb\) > .*debugger_test\.py\([0-9]+\).*
|
|
@jax\.jit
|
|
def g\(x\):
|
|
-> y = f\(x\)
|
|
return jnp\.exp\(y\)
|
|
.*
|
|
\(jdb\) array\(2\., dtype=float32\)
|
|
\(jdb\) > .*debugger_test\.py\([0-9]+\)
|
|
def f\(x\):
|
|
y = jnp\.sin\(x\)
|
|
-> debugger\.breakpoint\(stdin=stdin, stdout=stdout, backend="cli"\)
|
|
return y
|
|
.*
|
|
\(jdb\) """)
|
|
g(jnp.array(2., jnp.float32))
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_can_use_multiple_breakpoints(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p y", "c", "p y", "c"])
|
|
|
|
def f(x):
|
|
y = x + 1.
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, ordered=True,
|
|
backend="cli")
|
|
return y
|
|
|
|
@jax.jit
|
|
def g(x):
|
|
y = f(x) * 2.
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, ordered=True,
|
|
backend="cli")
|
|
return jnp.exp(y)
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
(jdb) array(3., dtype=float32)
|
|
(jdb) Entering jdb:
|
|
(jdb) array(6., dtype=float32)
|
|
(jdb) """)
|
|
g(jnp.array(2., jnp.float32))
|
|
jax.effects_barrier()
|
|
self.assertEqual(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_works_with_vmap(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p y", "c", "p y", "c"])
|
|
|
|
# On TPU, the breakpoints can be reordered inside of vmap but can be fixed
|
|
# by ordering sends.
|
|
# TODO(sharadmv): change back to ordered = False when sends are ordered
|
|
ordered = jax.default_backend() == "tpu"
|
|
|
|
def f(x):
|
|
y = x + 1.
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, ordered=ordered,
|
|
backend="cli")
|
|
return 2. * y
|
|
|
|
@jax.jit
|
|
@jax.vmap
|
|
def g(x):
|
|
y = f(x)
|
|
return jnp.exp(y)
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
(jdb) array(1., dtype=float32)
|
|
(jdb) Entering jdb:
|
|
(jdb) array(2., dtype=float32)
|
|
(jdb) """)
|
|
g(jnp.arange(2., dtype=jnp.float32))
|
|
jax.effects_barrier()
|
|
self.assertEqual(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_works_with_pmap(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
if jax.local_device_count() < 2:
|
|
raise unittest.SkipTest("Test requires >= 2 devices.")
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p y", "c", "p y", "c"])
|
|
|
|
def f(x):
|
|
y = jnp.sin(x)
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
|
|
@jax.pmap
|
|
def g(x):
|
|
y = f(x)
|
|
return jnp.exp(y)
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) array\(.*, dtype=float32\)
|
|
\(jdb\) Entering jdb:
|
|
\(jdb\) array\(.*, dtype=float32\)
|
|
\(jdb\) """)
|
|
g(jnp.arange(2., dtype=jnp.float32))
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_works_with_pjit(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
if jax.default_backend() != "tpu":
|
|
raise unittest.SkipTest("`pjit` doesn't work with CustomCall.")
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p y", "c"])
|
|
|
|
def f(x):
|
|
y = x + 1
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
return y
|
|
|
|
def g(x):
|
|
y = f(x)
|
|
return jnp.exp(y)
|
|
g = pjit.pjit(
|
|
g,
|
|
in_shardings=jax.sharding.PartitionSpec("dev"),
|
|
out_shardings=jax.sharding.PartitionSpec("dev"),
|
|
)
|
|
with jax.sharding.Mesh(np.array(jax.devices()), ["dev"]):
|
|
arr = (1 + np.arange(8)).astype(np.int32)
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) {}
|
|
\(jdb\) """.format(re.escape(repr(arr))))
|
|
g(jnp.arange(8, dtype=jnp.int32))
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_uses_local_before_global_scope(self):
|
|
stdin, stdout = make_fake_stdin_stdout(["p foo", "c"])
|
|
|
|
foo = "outer"
|
|
|
|
def f(x):
|
|
foo = "inner"
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
del foo
|
|
return x
|
|
|
|
del foo
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) 'inner'
|
|
\(jdb\) """)
|
|
f(2.)
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_debugger_accesses_globals(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p foo", "c"])
|
|
|
|
@jax.jit
|
|
def g():
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) \*\*\* NameError: name 'foo' is not defined
|
|
\(jdb\) """)
|
|
g()
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_can_limit_num_frames(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
stdin, stdout = make_fake_stdin_stdout(["u", "p x", "c"])
|
|
|
|
def g():
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli",
|
|
num_frames=2)
|
|
|
|
@jax.jit
|
|
def f():
|
|
x = 2
|
|
g()
|
|
return x
|
|
_ = f()
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) .*
|
|
.*
|
|
.*
|
|
.*
|
|
.*
|
|
.*
|
|
.*
|
|
\(jdb\) 2
|
|
\(jdb\) """)
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["u", "u", "c"])
|
|
|
|
def g2():
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli",
|
|
num_frames=2)
|
|
|
|
@jax.jit
|
|
def f2():
|
|
x = 2
|
|
g2()
|
|
return x
|
|
|
|
expected = ".*At topmost frame.*"
|
|
_ = f2()
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
def test_can_handle_dictionaries_with_unsortable_keys(self):
|
|
if xla_bridge.get_backend().runtime_type == 'stream_executor':
|
|
raise unittest.SkipTest('Host callback not supported for runtime type: stream_executor.')
|
|
|
|
stdin, stdout = make_fake_stdin_stdout(["p x", "p weird_dict",
|
|
"p weirder_dict", "c"])
|
|
|
|
@jax.jit
|
|
def f():
|
|
weird_dict = {(lambda x: x): 2., (lambda x: x * 2): 3}
|
|
weirder_dict = {(lambda x: x): weird_dict}
|
|
x = 2.
|
|
debugger.breakpoint(stdin=stdin, stdout=stdout, backend="cli")
|
|
del weirder_dict
|
|
return x
|
|
expected = _format_multiline(r"""
|
|
Entering jdb:
|
|
\(jdb\) 2.0
|
|
\(jdb\) <cant_flatten>
|
|
\(jdb\) <cant_flatten>
|
|
\(jdb\) """)
|
|
_ = f()
|
|
jax.effects_barrier()
|
|
self.assertRegex(stdout.getvalue(), expected)
|
|
|
|
if __name__ == '__main__':
|
|
absltest.main(testLoader=jtu.JaxTestLoader())
|