chore: import upstream snapshot with attribution

python/ray/train/v2/jax/__init__.py

@@ -0,0 +1,15 @@
+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"]