MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / jit_script_compat

Function jit_script_compat

deepspeed/utils/torch.py:49–73  ·  view source on GitHub ↗
(fn)

Source from the content-addressed store, hash-verified

47
48
49def jit_script_compat(fn):
50 fn_name = getattr(fn, "__qualname__", getattr(fn, "__name__", repr(fn)))
51
52 can_try_compile = (required_torch_version(min_version=2.0) and hasattr(torch, "compile")
53 and not (sys.version_info >= (3, 12) and not required_torch_version(min_version=2.4)))
54
55 if can_try_compile:
56 try:
57 return torch.compile(fn)
58 except Exception:
59 _logger.debug(
60 "torch.compile failed for %s, falling back to torch.jit.script",
61 fn_name,
62 exc_info=True,
63 )
64
65 try:
66 return torch.jit.script(fn)
67 except Exception:
68 _logger.debug(
69 "torch.jit.script failed for %s, returning unmodified function",
70 fn_name,
71 exc_info=True,
72 )
73 return fn

Callers

nothing calls this directly

Calls 2

required_torch_versionFunction · 0.85
compileMethod · 0.45

Tested by

no test coverage detected