return a + 1.0 - a
self.checkScript(test_rand, ())
+ fn = torch.jit.script(test_rand)
+ out = fn()
+ self.assertEqual(out.dtype, torch.double)
+ g = fn.graph_for()
+ # Testing shape analysis correctly setting type
+ FileCheck().check("Double(*, *)").check_not("Float(*, *)").run(g)
+
+ @torch.jit.script
+ def randint():
+ return torch.randint(0, 5, [1, 2])
+ out = randint()
+ self.assertEqual(out.dtype, torch.double)
+ # although the type should be int here, testing that the runtime dtype
+ # and shape analysis dtype is the same.
+ FileCheck().check("Double(*, *)").check_not("Float(*, *)").run(randint.graph_for())
def test_erase_number_types(self):
def func(a):
auto first_scalar_type = (*maybe_tensor_types)[0]->scalarType();
auto second_scalar_type = (*maybe_tensor_types)[1]->scalarType();
size_t arg_for_type = 0;
- if (c10::promoteTypes(first_scalar_type, second_scalar_type) != first_scalar_type) {
+ if (c10::promoteTypes(first_scalar_type, second_scalar_type) !=
+ first_scalar_type) {
arg_for_type = 1;
}
return {broadcast(*maybe_tensor_types, arg_for_type)};
if (!maybe_dtype_option)
return {};
auto dtype =
- (maybe_dtype_option->isNone() ? at::kFloat
+ (maybe_dtype_option->isNone() ? at::kDouble
: maybe_dtype_option->toScalarType());
return {DimensionedTensorType::create(dtype, device, dim)};