test fix

.github/workflows/build_and_test.yml

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

.github/workflows/build_upstream_image.yml

@@ -5,7 +5,7 @@ on:
       - master
       - r[0-9]+.[0-9]+
     paths-ignore:
-      - 'experimental/torch_xla2/**'
+      - 'experimental/**'
   workflow_dispatch:
 jobs:
   build:

experimental/torch_xla2/test/test_ops.py

@@ -7,6 +7,7 @@
     instantiate_device_type_tests, ops)
 from torch.utils import _pytree as pytree
 from torch_xla2 import tensor
+import torch_xla2
 
 
 skiplist = {

experimental/torch_xla2/torch_xla2/__init__.py

@@ -16,7 +16,6 @@
 
 from jax._src import xla_bridge
 os.environ.setdefault('ENABLE_RUNTIME_UPTIME_TELEMETRY', '1')
-jax.config.update('jax_enable_x64', True)
 
 # torch_xla2:oss-begin
 old_pjrt_options = jax.config.jax_pjrt_client_create_options
@@ -80,4 +79,16 @@ def disable_globally():
   unsupported_dtype=unsupported_dtype)
 
 import jax
-torch._register_device_module('jax', jax)
+torch._register_device_module('jax', jax)
+
+
+def enable_accuracy_mode():
+  jax.config.update('jax_enable_x64', True)
+  jax.config.update('jax_default_matmul_precision', 'highest')
+  default_env().config.internal_respect_torch_return_dtypes = True
+
+
+def enable_performance_mode():
+  jax.config.update('jax_enable_x64', False)
+  jax.config.update('jax_default_matmul_precision', 'default')
+  default_env().config.internal_respect_torch_return_dtypes = False

experimental/torch_xla2/torch_xla2/config.py

@@ -15,3 +15,4 @@ class Configuration:
     # device
     treat_cuda_as_jax_device: bool = True
     use_torch_native_for_cpu_tensor: bool = False
+    internal_respect_torch_return_dtypes: bool = False