Support for torch.float16

[maxfrei750]

Hi! When using dr.wrap_ad in the following way:

import drjit as dr
import torch


@dr.wrap_ad(source="drjit", target="torch")
def test():
    return torch.tensor(1, dtype=torch.float16)


a = test()

I get:

Exception has occurred: KeyError
torch.float16
  File "/src/mitsuba3/build/python/drjit/router.py", line 5801, in torch_to_drjit
    dtype_str = {
  File "/src/mitsuba3/build/python/drjit/router.py", line 5860, in eval
    return torch_to_drjit(self.res_torch)
  File "/src/mitsuba3/build/python/drjit/router.py", line 5612, in custom
    output = inst.eval(*_dr.detach(kwargs['args']))
  File "/src/mitsuba3/build/python/drjit/router.py", line 5873, in f
    return _dr.custom(ToTorch, args)
  File "/workspace/development/mwe.py", line 10, in <module>
    a = test()

Is this expected behavior, i.e. mitsuba does not support float16? Or is it just missing in the dictionary here:

https://github.com/mitsuba-renderer/drjit/blob/fdb652f90e5fee40621090742b5d6bb326acfae4/drjit/router.py#L5801-5809

Please excuse me, if this is being mentioned somewhere in the documentation, but a quick search did not reveal anything. Searching issues and discussion for 'float16' also did not yield anything.

[Speierers]

Hi @maxfrei750,

Half precision floats are not really supported in Mitsuba 3 and Dr.Jit at this point. You will need to use float32 on the PyTorch side, or perform the conversion yourself before passing the tensor to dr.wrap_ad().

In the future we could look into adding support for float16 types in Dr.Jit but this isn't currently on our roadmap.

[maxfrei750]

Thanks for the clarification. Indeed, calling .float() on the returned tensor solves the problem.