Pylint error

Author: NeilGirdhar

src/array_api_extra/_delegation.py

@@ -174,16 +174,16 @@ def one_hot(
     # Delegate where possible.
     if is_jax_namespace(xp):
         assert is_jax_array(x)
-        from jax.nn import one_hot
+        from jax.nn import one_hot as jax_one_hot
 
-        return one_hot(x, num_classes, dtype=dtype, axis=axis)
+        return jax_one_hot(x, num_classes, dtype=dtype, axis=axis)
     if is_torch_namespace(xp):
         assert is_torch_array(x)
-        from torch.nn.functional import one_hot
+        from torch.nn.functional import one_hot as torch_one_hot
 
         x = xp.astype(x, xp.int64)  # PyTorch only supports int64 here.
         try:
-            out = one_hot(x, num_classes)
+            out = torch_one_hot(x, num_classes)
         except RuntimeError as e:
             raise IndexError from e
         out = xp.astype(out, dtype)