Update

chaobrain · Jun 14, 2024 · c172528 · c172528
1 parent 58c9bbf
commit c172528
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/brainunit/math/_compat_numpy_funcs_remove_unit.py b/brainunit/math/_compat_numpy_funcs_remove_unit.py
@@ -83,9 +83,9 @@ def sign(x: Union[Array, Quantity]) -> Array:
 def histogram(
     x: Union[Array, Quantity],
     bins: jax.typing.ArrayLike = 10,
-    range: Sequence[jax.typing.ArrayLike] | None = None,
-    weights: jax.typing.ArrayLike | None = None,
-    density: bool | None = None
+    range: Optional[Sequence[jax.typing.ArrayLike]] = None,
+    weights: Optional[jax.typing.ArrayLike] = None,
+    density: Optional[bool] = None
 ) -> tuple[Array, Array]:
   """
   Compute the histogram of a set of data.
@@ -280,9 +280,9 @@ def cov(
     y: Optional[Union[Array, Quantity]] = None,
     rowvar: bool = True,
     bias: bool = False,
-    ddof: int | None = None,
-    fweights: jax.typing.ArrayLike | None = None,
-    aweights: jax.typing.ArrayLike | None = None
+    ddof: Optional[int] = None,
+    fweights: Optional[jax.typing.ArrayLike] = None,
+    aweights: Optional[jax.typing.ArrayLike] = None
 ) -> Array:
   """
   Estimate a covariance matrix, given data and weights.

diff --git a/brainunit/math/_compat_numpy_linear_algebra.py b/brainunit/math/_compat_numpy_linear_algebra.py
@@ -40,7 +40,7 @@ def dot(
     b: Union[Array, Quantity],
     *,
     precision: Any = None,
-    preferred_element_type: jax.typing.DTypeLike | None = None
+    preferred_element_type: Optional[jax.typing.DTypeLike] = None
 ) -> Union[Array, Quantity]:
   """
   Dot product of two arrays or quantities.
@@ -80,7 +80,7 @@ def vdot(
     b: Union[Array, Quantity],
     *,
     precision: Any = None,
-    preferred_element_type: jax.typing.DTypeLike | None = None
+    preferred_element_type: Optional[jax.typing.DTypeLike] = None
 ) -> Union[Array, Quantity]:
   """
   Perform a conjugate multiplication of two 1D vectors.
@@ -120,7 +120,7 @@ def inner(
     b: Union[Array, Quantity],
     *,
     precision: Any = None,
-    preferred_element_type: jax.typing.DTypeLike | None = None
+    preferred_element_type: Optional[jax.typing.DTypeLike] = None
 ) -> Union[Array, Quantity]:
   """
   Inner product of two arrays or quantities.
@@ -220,7 +220,7 @@ def matmul(
     b: Union[Array, Quantity],
     *,
     precision: Any = None,
-    preferred_element_type: jax.typing.DTypeLike | None = None
+    preferred_element_type: Optional[jax.typing.DTypeLike] = None
 ) -> Union[Array, Quantity]:
   """
   Matrix product of two arrays or quantities.
@@ -260,7 +260,7 @@ def trace(
     offset: int = 0,
     axis1: int = 0,
     axis2: int = 1,
-    dtype: jax.typing.DTypeLike | None = None,
+    dtype: Optional[jax.typing.DTypeLike] = None,
     out: None = None
 ) -> Union[Array, Quantity]:
   """

diff --git a/brainunit/math/_compat_numpy_misc.py b/brainunit/math/_compat_numpy_misc.py
@@ -98,7 +98,7 @@ def einsum(
     out: None = None,
     optimize: Union[str, bool] = "optimal",
     precision: jax.lax.PrecisionLike = None,
-    preferred_element_type: Union[jax.typing.DTypeLike, None] = None,
+    preferred_element_type: Optional[jax.typing.DTypeLike] = None,
     _dot_general: Callable[..., jax.Array] = jax.lax.dot_general,
 ) -> Union[jax.Array, Quantity]:
   """