mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Internal change

Browse Source 
PiperOrigin-RevId: 464052046
This commit is contained in:
jax authors 2022-07-29 04:46:33 -07:00
parent 7b17ee87a9
commit 80eb641f6a
1 changed files with 7 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
jax/_src/debugger/colab_lib.py
View File

@ -17,20 +17,19 @@ from __future__ import annotations
import abc
import dataclasses
import functools
import sys
import uuid
from typing import Any, Dict, List, Union
# pylint: disable=g-import-not-at-top
# pytype: disable=import-error
try:
IS_COLAB_ENABLED = "google.colab" in sys.modules
if IS_COLAB_ENABLED:
# pylint: disable=g-import-not-at-top
# pytype: disable=import-error
from google.colab import output
from IPython import display
IS_COLAB_ENABLED = True
except ImportError:
IS_COLAB_ENABLED = False
# pytype: enable=import-error
# pylint: enable=g-import-not-at-top
# pytype: enable=import-error
# pylint: enable=g-import-not-at-top
class DOMElement(metaclass=abc.ABCMeta):