I run into additional type checking frequently when I try to write a function that uses these math functions. For example: let’s say I want to have one utility function that computes Euclidean distance. Across my whole project, I might use it on tensors or Python scalars.
For Python scalars, I have to use something like:
import math
def euclidean_distance(a, b):
    """a and b are either Python scalars or tensors with size 1"""
    return math.sqrt(a**2 + b**2)

This will work for all Python scalars and all torch tensors of size 1. But for PyTorch tensors, generally, I have to use:
import torch
def euclidean_distance(a, b):
    """a and b are arbitrary tensors"""
    return torch.sqrt(a**2 + b**2)

Is there some way to allow for:
import torch
def euclidean_distance(a, b):
    """a and b are arbitrary tensors or Python scalars"""
    return torch.sqrt(a**2 + b**2)

without uses having to check like:
import torch
def euclidean_distance(a, b):
    """a and b are arbitrary tensors or Python scalars"""
    # Handle tensor vs not tensor inputs
    a_not_tensor = False
    if not torch.is_tensor(a):
        a_not_tensor = True
        a = torch.tensor([a])
    b_not_tensor = False
    if not torch.is_tensor(b):
        b_not_tensor = True
        b = torch.tensor([b])

    # Compute the distance
    dist = torch.sqrt(a**2 + b**2)

    # If both inputs are Python scalars, return a Python scalar
    if a_not_tensor and b_not_tensor:
        return dist.item()

    # If either input was a tensor, return a tensor
    return dist

Is there a reason this functionality doesn’t exist? It’s not limited to torch.sqrt(), but all PyTorch math functions. I imagine it’s easier to have torch.sqrt() accept Python scalars then to automatically vectorize math.sqrt() (a la Python build-ins like abs()), but I suppose either would solve the problem.
[An aside: this is kind of related to the lack of a torch.pi constant (1 1, 2 1, 3). Requiring import math for certain things is fine to remove redundant constants, etc, but it would also be useful to provide the cross-functionality when needed.]