moving torch_xla2 experimental code in-tree (#6463)

.github/workflows/build_and_test.yml

@@ -3,10 +3,14 @@ on:
     branches:
       - master
       - r[0-9]+.[0-9]+
+    paths-ignore:
+      - 'experimental/torch_xla2/**'
   push:
     branches:
       - master
       - r[0-9]+.[0-9]+
+    paths-ignore:
+      - 'experimental/torch_xla2/**'
   workflow_dispatch:
 
 concurrency:

experimental/torch_xla2/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2023, pytorch-tpu
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

experimental/torch_xla2/README.md

@@ -0,0 +1,3 @@
+# torchxla2
+
+This directory contains things that are in the top-level git repository

experimental/torch_xla2/dev-requirements.txt

@@ -0,0 +1,8 @@
+-r requirements.txt
+pytest
+yapf
+tabulate
+transformers
+tf-nightly
+--pre -f https://download.pytorch.org/whl/nightly/torch_nightly.html
+torchvision

experimental/torch_xla2/format.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+set -ex
+
+yapf --recursive -i *.py test torch_xla2

experimental/torch_xla2/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+
+[project]
+version = "0.0.1"
+name = "torch_xla2"
+dependencies = [
+    "torch>=2.1",
+    "jax>=0.4.24",
+    "jaxlib",
+]
+
+requires-python = ">=3.10"
+license = {file = "LICENSE"}

experimental/torch_xla2/requirements.txt

@@ -0,0 +1,5 @@
+jax==0.4.24.dev20240202
+jaxlib==0.4.24.dev20240202
+numpy==1.26.3
+torch==2.2.0
+typing_extensions==4.9.0

experimental/torch_xla2/test/BUILD

@@ -0,0 +1,31 @@
+# TODO(hanq): describe this package.
+
+load(
+    "//third_party/py/torch/google/bazel_rules/rules_python/python:defs.bzl",
+    "py_library",
+    "py_test",
+)
+
+package(
+    default_applicable_licenses = ["//devtools/compliance/licenses:no_external_contributions"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "test_base",
+    srcs = ["test_base.py"],
+    deps = [
+        "//testing/pybase",
+    ],
+)
+
+py_test(
+    name = "test_core_aten_ops",
+    srcs = ["test_core_aten_ops.py"],
+    deps = [
+        ":test_base",
+        "//third_party/py/absl:app",
+        "//third_party/py/torch/google/_torx",
+    ],
+)

experimental/torch_xla2/test/llama/BUILD

@@ -0,0 +1,25 @@
+# TODO(hanq): describe this package.
+load(
+    "//third_party/py/torch/google/bazel_rules/rules_python/python:defs.bzl",
+    "py_test",
+)
+
+package(
+    default_applicable_licenses = ["//devtools/compliance/licenses:no_external_contributions"],
+    default_visibility = ["//visibility:public"],
+    licenses = ["notice"],
+)
+
+py_test(
+    name = "test_llama",
+    srcs = [
+        "llama_model.py",
+        "test_llama.py",
+    ],
+    deps = [
+        "//third_party/py/jax",
+        "//third_party/py/torch:pytorch",
+        "//third_party/py/torch/google/_torx",
+        "//third_party/py/torch/google/_torx/test:test_base",
+    ],
+)