[array API] clean up some unused/unnecessary code

jax/experimental/array_api/__init__.py

@@ -155,6 +155,7 @@
     square as square,
     squeeze as squeeze,
     stack as stack,
+    std as std,
     subtract as subtract,
     sum as sum,
     take as take,
@@ -174,6 +175,7 @@
     unique_inverse as unique_inverse,
     unique_values as unique_values,
     unstack as unstack,
+    var as var,
     vecdot as vecdot,
     where as where,
     zeros as zeros,
@@ -199,11 +201,6 @@
     hypot as hypot,
 )
 
-from jax.experimental.array_api._statistical_functions import (
-    std as std,
-    var as var,
-)
-
 from jax.experimental.array_api._utility_functions import (
     __array_namespace_info__ as __array_namespace_info__,
 )

jax/experimental/array_api/_array_methods.py

@@ -14,14 +14,9 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 import jax
 from jax._src.array import ArrayImpl
 from jax.experimental.array_api._version import __array_api_version__
-from jax.sharding import Sharding
-
-from jax._src.lib import xla_extension as xe
 
 
 def _array_namespace(self, /, *, api_version: None | str = None):

jax/experimental/array_api/linalg.py

@@ -31,12 +31,11 @@
   svd as svd,
   svdvals as svdvals,
   tensordot as tensordot,
+  trace as trace,
   vecdot as vecdot,
   vector_norm as vector_norm,
 )
 
-from jax.numpy.linalg import trace as trace
-
 from jax.experimental.array_api._linear_algebra_functions import (
   matrix_rank as matrix_rank,
   pinv as pinv,