array_api_compat_cross_library_function_numpy_pytorch_cupy.py

1import array_api_compat
2from array_api_compat import array_namespace
3
4def standardized_function(x):
5    # array_namespace(x) returns the array API compatible namespace 
6    # for the library that created x (e.g., numpy, torch, cupy)
7    xp = array_namespace(x)
8    
9    # Use the namespace xp to call functions defined in the Array API standard
10    return xp.mean(x, axis=0)
11
12# Example usage with NumPy
13import numpy as np
14x_np = np.array([1.0, 2.0, 3.0])
15print(f"NumPy result: {standardized_function(x_np)}")
16
17# Example usage with PyTorch (if installed)
18try:
19    import torch
20    x_torch = torch.tensor([1.0, 2.0, 3.0])
21    print(f"PyTorch result: {standardized_function(x_torch)}")
22except ImportError:
23    pass