Add bool tensor support for where

[H-Huang]

Stack from ghstack:

• Enable diag for bool Tensors #47455 Enable diag for bool Tensors
• Add bool tensor support for where #47454 Add bool tensor support for where

Differential Revision: D24772482

---

Description

Fixes #26247
This now works:

a = torch.tensor([True, False], device='cpu')
x = torch.tensor([True, True], device='cpu')
y = torch.tensor([False, False], device='cpu')
res = torch.where(a, x, y)

a = torch.tensor([True, False], device='cuda')
x = torch.tensor([True, True], device='cuda')
y = torch.tensor([False, False], device='cuda')
res = torch.where(a, x, y)

Also updated where_tensor test to include bools

Test:
pytest test/test_torch.py -k test_where_tensor -vs

[dr-ci]

💊 CI failures summary and remediations

As of commit 147bed7 (more details on the Dr. CI page):

---

• 1/1 failures introduced in this PR

---

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

pytorch_windows_vs2019_py36_cuda10.1_test2 (1/1)

Step: "Test" (full log | diagnosis details | 🔁 rerun)

distributed\_pipeline\sync\test_balance.py::test_balance_by_time_loop_resets_input FAILED [ 38%]

rootdir: C:\Users\circleci\project 
plugins: hypothesis-4.53.2 
collecting ... collected 18 items 
 
distributed\_pipeline\sync\test_balance.py::test_blockpartition PASSED   [  5%] 
distributed\_pipeline\sync\test_balance.py::test_blockpartition_zeros PASSED [ 11%] 
distributed\_pipeline\sync\test_balance.py::test_blockpartition_non_positive_partitions PASSED [ 16%] 
distributed\_pipeline\sync\test_balance.py::test_blockpartition_short_sequence PASSED [ 22%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_time[cpu] SKIPPED [ 27%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_time[cuda] SKIPPED [ 33%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_time_loop_resets_input FAILED [ 38%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_size_latent PASSED [ 44%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_size_param PASSED [ 50%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_size_param_scale PASSED [ 55%] 
distributed\_pipeline\sync\test_balance.py::test_layerwise_sandbox[cpu] PASSED [ 61%] 
distributed\_pipeline\sync\test_balance.py::test_layerwise_sandbox[cuda] PASSED [ 66%] 
distributed\_pipeline\sync\test_balance.py::test_sandbox_during_profiling[cpu] PASSED [ 72%] 
distributed\_pipeline\sync\test_balance.py::test_sandbox_during_profiling[cuda] PASSED [ 77%] 
distributed\_pipeline\sync\test_balance.py::test_not_training PASSED     [ 83%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_time_tuple PASSED [ 88%] 
distributed\_pipeline\sync\test_balance.py::test_balance_by_size_tuple PASSED [ 94%] 

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 5 times.

[codecov]

Codecov Report

Merging #47454 into gh/H-Huang/1/base will increase coverage by 0.00%.
The diff coverage is n/a.

@@                Coverage Diff                 @@
##           gh/H-Huang/1/base   #47454   +/-   ##
==================================================
  Coverage              81.44%   81.44%           
==================================================
  Files                   1798     1798           
  Lines                 188223   188223           
==================================================
+ Hits                  153293   153303   +10     
+ Misses                 34930    34920   -10     

[facebook-github-bot]

@H-Huang merged this pull request in 3253ccb.