Pred type causes issues about xla

[dlibenzi]

@ailzhang Is this TPU or CPU?

[dlibenzi]

Must be TPU as this works on CPU:

  def test_pred_type(self):
    xla_device = xm.xla_device()
    a = (torch.rand(4).to(xla_device) >= 0.5)
    b = (torch.rand(4).to(xla_device) >= 0.5)
    c = torch.cat([a, b], dim=0)
    d = c.sum()
    print(c)
    print(d)

tensor([0, 0, 0, 0, 0, 1, 0, 0], dtype=torch.uint8)
tensor(1, dtype=torch.uint8)

[dlibenzi]

Yep, it is:

RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:242 : Check failed: mwait.Wait() == ::tensorflow::Status::OK() (Internal: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:222 : Check failed: session->session()->Run( session_work.feed_inputs, session_work.outputs_handles, &outputs) == ::tensorflow::Status::OK() (Invalid argument: Invalid HLO for TPU: U8 and S8 types are not supported: %x.13 = u8[] parameter(0)
	 [[{{node XRTCompile}}]]
	 [[XRTCompile_G3]] vs. OK)