import torch
import functorch

# Reduced from test_torch.py: test_exponential
def poc7():
  device = 'cpu'

  test = (-0, float('inf'))
  t = torch.empty((1,), device=device, dtype=torch.bfloat16).exponential_(test[0])
  print(t.item() == test[1])


print("CPU")
poc7()
print()

print("CPU Functionalize")
functorch.functionalize(poc7)()

CPU
True

CPU Functionalize
/workspaces/work/pytorch/torch/_functorch/deprecated.py:93: UserWarning: We've integrated functorch into PyTorch. As the final step of the integration, functorch.functionalize is deprecated as of PyTorch 2.0 and will be deleted in a future version of PyTorch >= 2.3. Please use torch.func.functionalize instead; see the PyTorch 2.0 release notes and/or the torch.func migration guide for more details https://pytorch.org/docs/master/func.migrating.html
  warn_deprecated('functionalize')
Traceback (most recent call last):
  File "poc7.py", line 19, in <module>
    functorch.functionalize(poc7)()
  File "/workspaces/work/pytorch/torch/_functorch/vmap.py", line 40, in fn
    return f(*args, **kwargs)
  File "/workspaces/work/pytorch/torch/_functorch/eager_transforms.py", line 1582, in wrapped
    func_outputs = func(*func_args, **func_kwargs)
  File "poc7.py", line 10, in poc7
    t = torch.empty((1,), device=device, dtype=torch.bfloat16).exponential_(test[0])
  File "/workspaces/work/pytorch/torch/_refs/__init__.py", line 432, in _fn
    return fn(a, *args, out=a, **kwargs)
  File "/workspaces/work/pytorch/torch/_prims_common/wrappers.py", line 223, in _fn
    result = fn(*args, **kwargs)
  File "/workspaces/work/pytorch/torch/_prims_common/wrappers.py", line 133, in _fn
    result = fn(**bound.arguments)
  File "/workspaces/work/pytorch/torch/_refs/__init__.py", line 5190, in exponential
    return -1 / rate * torch.log1p(-torch.rand_like(self))
ZeroDivisionError: float division by zero