wehub-resource-sync
16 lines
483 B
Python
16 lines
483 B
Python
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
try:
|
|
import jax # noqa: F401
|
|
except ModuleNotFoundError as exception:
|
|
raise ModuleNotFoundError(
|
|
"Jax isn't installed. To install Jax, please check"
|
|
" `https://github.com/google/jax#installation` for the instructions."
|
|
) from exception
|
|
|
|
from ray.train.v2.jax.config import JaxConfig
|
|
from ray.train.v2.jax.jax_trainer import JaxTrainer
|
|
|
|
__all__ = ["JaxConfig", "JaxTrainer"]
|