Add support for XLA_DISABLE_FUNCTIONALIZATION flag

[wonjoolee95]

Add support for XLA_DISABLE_FUNCTIONALIZATION flag

This change adds support for XLA_DISABLE_FUNCTIONALIZATION. This involves restoring some of the deleted ops when we enabled functionalization and updating some ops to fall back to pre-functionalization state. This change also adds a new test suite that runs test_operations.py with functionalization disabled.

Test plan: source run_tests.sh and CI

---

Ops that are completely restored:

• as_strided
• as_strided_
• diagonal
• expand
• view

Ops that are modified by checking XLA_DISABLE_FUNCTIONALIZATION flag

• convolution_backward
• new_empty_strided_symint
• select_symint
• slice
• as_strided

---

Skipped tests (these two tests were added as part of the funtionalization commit, skipping when XLA_DISABLE_FUNCTIONALIZATION is set to true)

• test_replace_xla_tensor
• test_set

Fixed tests

• test_ailing_slice (slice)
• test_as_strided (as_strided, as_strided_)
• test_conv2d_backward (convlution_backward)
• test_diagonal (diagonal)
• test_expand (expand)
• test_flip (view)
• test_index_bool (select_symint)
• test_inplace_view (view, new_empty_strided_symint)
• test_matmul_integer_types (view)
• test_negative_slice (slice)
• test_pred_one_hot (select_symint)
• test_rrelu_module (select_symint)
• test_slice (slice)
• test_spooky_ailing (select_symint)
• test_squeeze_nonzero (new_empty_strided_symint)
• test_view_1718 (view)
• test_view_and_copy_ (view)
• test_view_out_computation (view)
• TestParallelTensorMNIST.test (convlution_backward)
• TestParallelTensorResnet18.test (convlution_backward)

cc @alanwaketan

[wonjoolee95]

Locally I can see that test_operations.py succeeds when XLA_DISABLE_FUNCTIONALIZATION=1:

----------------------------------------------------------------------
Ran 151 tests in 170.539s

OK (skipped=3)
(base) jenkins@26d7adccbc26:/workspace/pytorch/xla/test$ 

And also succeeds without the flag:

----------------------------------------------------------------------
Ran 151 tests in 130.983s

OK (skipped=1)
(base) jenkins@26d7adccbc26:/workspace/pytorch/xla/test$ 

I'll let the CI verify other tests.

[alanwaketan]

Generally LGTM. Please address my comments. Thanks for coming out with this PR so quickly.

[alanwaketan]

Instead of skipping the tests here, maybe you can utilize the env and have different assertEquals.

[wonjoolee95]

These two tests didn't exist before functionalization was enabled (they were newly added as of #4158), should we still even run these when functionalization is disabled?

[alanwaketan]

Forgot the details. Thanks for pointing it out. Then, let's skip.

[wonjoolee95]

Thanks for the review, Jiewen. Merging this now as CIs are green.