openscvx 0.3.2.dev316__tar.gz → 0.3.2.dev317__tar.gz

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/.github/workflows/_docs.yml

@@ -60,7 +60,7 @@ jobs:
 
     - name: Build documentation (check only)
       if: ${{ !inputs.version }}
-      run: mkdocs build
+      run: mkdocs build --strict
 
     - name: Fetch gh-pages branch
       if: inputs.version

{openscvx-0.3.2.dev316/openscvx.egg-info → openscvx-0.3.2.dev317}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.3.2.dev316
+Version: 0.3.2.dev317
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.2.dev316'
-__version_tuple__ = version_tuple = (0, 3, 2, 'dev316')
+__version__ = version = '0.3.2.dev317'
+__version_tuple__ = version_tuple = (0, 3, 2, 'dev317')
 
-__commit_id__ = commit_id = 'ga36b834de'
+__commit_id__ = commit_id = 'g443a82694'

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/algorithms/optimization_results.py

@@ -15,36 +15,33 @@ class OptimizationResults:
     storage for plotting and application-specific data.
 
     Attributes:
-        converged (bool): Whether the optimization successfully converged
-        t_final (float): Final time of the optimized trajectory
+        converged (bool): Whether the optimization successfully converged.
+        t_final (float): Final time of the optimized trajectory.
         x_guess (np.ndarray): Optimized state trajectory at discretization nodes,
-            shape (N, n_states)
+            shape (N, n_states).
         u_guess (np.ndarray): Optimized control trajectory at discretization nodes,
-            shape (N, n_controls)
-
-        # Dictionary-based Access
+            shape (N, n_controls).
         nodes (dict[str, np.ndarray]): Dictionary mapping state/control names to arrays
             at optimization nodes. Includes both user-defined and augmented variables.
         trajectory (dict[str, np.ndarray]): Dictionary mapping state/control names to arrays
             along the propagated trajectory. Added by post_process().
-
-        # SCP Iteration History (for convergence analysis)
-        x_history (list[np.ndarray]): State trajectories from each SCP iteration
-        u_history (list[np.ndarray]): Control trajectories from each SCP iteration
-        discretization_history (list[np.ndarray]): Time discretization from each iteration
-        J_tr_history (list[np.ndarray]): Trust region cost history
-        J_vb_history (list[np.ndarray]): Virtual buffer cost history
-        J_vc_history (list[np.ndarray]): Virtual control cost history
-
-        # Post-processing Results (added by propagate_trajectory_results)
-        t_full (Optional[np.ndarray]): Full time grid for interpolated trajectory
-        x_full (Optional[np.ndarray]): Interpolated state trajectory on full time grid
-        u_full (Optional[np.ndarray]): Interpolated control trajectory on full time grid
-        cost (Optional[float]): Total cost of the optimized trajectory
-        ctcs_violation (Optional[np.ndarray]): Continuous-time constraint violations
-
-        # User-defined Data
-        plotting_data (dict[str, Any]): Flexible storage for plotting and application data
+        x_history (list[np.ndarray]): State trajectories from each SCP iteration.
+        u_history (list[np.ndarray]): Control trajectories from each SCP iteration.
+        discretization_history (list[np.ndarray]): Time discretization from each iteration.
+        J_tr_history (list[np.ndarray]): Trust region cost history.
+        J_vb_history (list[np.ndarray]): Virtual buffer cost history.
+        J_vc_history (list[np.ndarray]): Virtual control cost history.
+        t_full (Optional[np.ndarray]): Full time grid for interpolated trajectory.
+            Added by propagate_trajectory_results.
+        x_full (Optional[np.ndarray]): Interpolated state trajectory on full time grid.
+            Added by propagate_trajectory_results.
+        u_full (Optional[np.ndarray]): Interpolated control trajectory on full time grid.
+            Added by propagate_trajectory_results.
+        cost (Optional[float]): Total cost of the optimized trajectory.
+            Added by propagate_trajectory_results.
+        ctcs_violation (Optional[np.ndarray]): Continuous-time constraint violations.
+            Added by propagate_trajectory_results.
+        plotting_data (dict[str, Any]): Flexible storage for plotting and application data.
     """
 
     # Core optimization results
@@ -105,7 +102,7 @@ class OptimizationResults:
         """Initialize the results object."""
         pass
 
-    def update_plotting_data(self, **kwargs):
+    def update_plotting_data(self, **kwargs: Any) -> None:
         """
         Update the plotting data with additional information.
 

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/config.py

@@ -211,14 +211,6 @@ class SimConfig:
             u (Control): Control object, must have .min and .max attributes for
                 bounds.
             total_time (float): The total simulation time.
-            idx_x_true (slice): Slice for true state indices.
-            idx_x_true_prop (slice): Slice for true propagation state indices.
-            idx_u_true (slice): Slice for true control indices.
-            idx_t (slice): Slice for time index.
-            idx_y (slice): Slice for constraint violation indices.
-            idx_y_prop (slice): Slice for propagation constraint violation
-                indices.
-            idx_s (slice): Slice for time dilation index.
             save_compiled (bool): If True, save and reuse compiled solver
                 functions. Defaults to False.
             ctcs_node_intervals (list, optional): Node intervals for CTCS

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/discretization/discretization.py

@@ -43,7 +43,13 @@ def dVdt(
         n_u (int): Number of controls.
         N (int): Number of nodes in trajectory.
         dis_type (str): Discretization type ("ZOH" or "FOH").
-        **params: Additional parameters passed to state_dot, A, and B.
+        S_x: State scaling matrix.
+        c_x: State offset vector.
+        S_u: Control scaling matrix.
+        c_u: Control offset vector.
+        inv_S_x: Inverse state scaling matrix.
+        inv_S_u: Inverse control scaling matrix.
+        params: Additional parameters passed to state_dot, A, and B.
 
     Returns:
         jnp.ndarray: Time derivative of augmented state vector.
@@ -117,14 +123,14 @@ def dVdt(
 
 
 def calculate_discretization(
-    x,
-    u,
+    x: np.ndarray,
+    u: np.ndarray,
     state_dot: callable,
     A: callable,
     B: callable,
     settings: Config,
     params: dict,
-):
+) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
     """Calculate the discretized system matrices.
 
     This function computes the discretized system matrices (A_bar, B_bar, C_bar)
@@ -137,20 +143,15 @@ def calculate_discretization(
         A (callable): Function computing state Jacobian.
         B (callable): Function computing control Jacobian.
         settings: Configuration settings for OpenSCvx.
-        custom_integrator (bool): Whether to use custom RK45 integrator.
-        debug (bool): Whether to use debug mode.
-        solver (str): Name of the solver to use.
-        rtol (float): Relative tolerance for integration.
-        atol (float): Absolute tolerance for integration.
-        dis_type (str): Discretization type ("ZOH" or "FOH").
-        **kwargs: Additional parameters passed to state_dot, A, and B.
+        params: Additional parameters passed to state_dot, A, and B.
 
     Returns:
-        tuple: (A_bar, B_bar, C_bar, z_bar, Vmulti) where:
+        tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+            (A_bar, B_bar, C_bar, x_prop, Vmulti) where:
             - A_bar: Discretized state transition matrix
             - B_bar: Discretized control influence matrix
             - C_bar: Discretized control influence matrix for next node
-            - z_bar: Defect vector
+            - x_prop: Propagated state
             - Vmulti: Full augmented state trajectory
     """
     # Unpack settings
@@ -232,7 +233,7 @@ def calculate_discretization(
     return A_bar, B_bar, C_bar, x_prop, Vmulti
 
 
-def get_discretization_solver(dyn: Dynamics, settings: Config):
+def get_discretization_solver(dyn: Dynamics, settings: Config) -> callable:
     """Create a discretization solver function.
 
     This function creates a solver that computes the discretized system matrices
@@ -240,7 +241,7 @@ def get_discretization_solver(dyn: Dynamics, settings: Config):
 
     Args:
         dyn (Dynamics): System dynamics object.
-        settings: Configuration settings for discretization.
+        settings (Config): Configuration settings for discretization.
 
     Returns:
         callable: A function that computes the discretized system matrices.

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/integrators/runge_kutta.py

@@ -63,7 +63,7 @@ def rk45_step(
     t: jnp.ndarray,
     y: jnp.ndarray,
     h: float,
-    *args
+    *args: Any
 ) -> jnp.ndarray:
     """
     Perform a single RK45 (Runge-Kutta-Fehlberg) integration step.
@@ -96,11 +96,11 @@ def solve_ivp_rk45(
     f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
     tau_final: float,
     y_0: jnp.ndarray,
-    args,
+    args: tuple,
     tau_0: float = 0.0,
     num_substeps: int = 50,
     is_not_compiled: bool = False,
-):
+) -> jnp.ndarray:
     """
     Solve an initial-value ODE problem using fixed-step RK45 integration.
 
@@ -145,14 +145,14 @@ def solve_ivp_diffrax(
     f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
     tau_final: float,
     y_0: jnp.ndarray,
-    args,
+    args: tuple,
     tau_0: float = 0.0,
     num_substeps: int = 50,
     solver_name: str = "Dopri8",
     rtol: float = 1e-3,
     atol: float = 1e-6,
-    extra_kwargs=None,
-):
+    extra_kwargs: dict = None,
+) -> jnp.ndarray:
     """
     Solve an initial-value ODE problem using a Diffrax adaptive solver.
 
@@ -207,16 +207,16 @@ def solve_ivp_diffrax_prop(
     f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
     tau_final: float,
     y_0: jnp.ndarray,
-    args,
+    args: tuple,
     tau_0: float = 0.0,
     num_substeps: int = 50,
     solver_name: str = "Dopri8",
     rtol: float = 1e-3,
     atol: float = 1e-6,
-    extra_kwargs=None,
+    extra_kwargs: dict = None,
     save_time: jnp.ndarray = None,
     mask: jnp.ndarray = None,
-):
+) -> jnp.ndarray:
     """
     Solve an initial-value ODE problem using a Diffrax adaptive solver.
     This function is specifically designed for use in the context of

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/lowered/unified.py

@@ -161,7 +161,7 @@ class UnifiedState:
         initial=0.0,
         final=0.0,
         augmented=False,
-    ):
+    ) -> None:
         """Append another state or create a new state variable.
 
         This method allows dynamic extension of the unified state, either by appending
@@ -536,7 +536,7 @@ class UnifiedControl:
         max=np.inf,
         guess=0.0,
         augmented=False,
-    ):
+    ) -> None:
         """Append another control or create a new control variable.
 
         This method allows dynamic extension of the unified control, either by appending

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/plotting/plotting.py

@@ -473,7 +473,7 @@ def plot_projections_2d(
     var_name: str = "position",
     velocity_var_name: str | None = None,
     cmap: str = "viridis",
-):
+) -> go.Figure:
     """Plot XY, XZ, YZ projections of a 3D variable.
 
     Useful for visualizing 3D trajectories in 2D plane views.
@@ -628,7 +628,7 @@ def plot_vector_norm(
     result: OptimizationResults,
     var_name: str,
     bounds: tuple[float, float] | None = None,
-):
+) -> go.Figure:
     """Plot the 2-norm of a vector variable over time.
 
     Useful for visualizing thrust magnitude, velocity magnitude, etc.

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/plotting/scp_iteration.py

@@ -13,7 +13,7 @@ def plot_scp_iterations(
     control_names: list[str] | None = None,
     cmap_name: str = "viridis",
     show_propagation: bool = True,
-):
+) -> go.Figure:
     """Plot all SCP iterations overlaid with colormap-based coloring.
 
     Shows the evolution of states and controls across SCP iterations. Early

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/problem.py

@@ -79,10 +79,10 @@ class Problem:
         dynamics_prop: Optional[dict] = None,
         states_prop: Optional[List[State]] = None,
         algebraic_prop: Optional[dict] = None,
-        licq_min=0.0,
-        licq_max=1e-4,
-        time_dilation_factor_min=0.3,
-        time_dilation_factor_max=3.0,
+        licq_min: float = 0.0,
+        licq_max: float = 1e-4,
+        time_dilation_factor_min: float = 0.3,
+        time_dilation_factor_max: float = 3.0,
         byof: Optional[ByofSpec] = None,
     ):
         """The primary class in charge of compiling and exporting the solvers.
@@ -108,15 +108,15 @@ class Problem:
                 Only specify additional states beyond optimization states. Used with dynamics_prop.
             algebraic_prop (dict, optional): Dictionary mapping names to symbolic expressions
                 for outputs evaluated (not integrated) during propagation.
-            licq_min: Minimum LICQ constraint value
-            licq_max: Maximum LICQ constraint value
-            time_dilation_factor_min: Minimum time dilation factor
-            time_dilation_factor_max: Maximum time dilation factor
-            byof: Expert mode only. Raw JAX functions to bypass symbolic layer.
-                See :class:`openscvx.expert.ByofSpec` for detailed documentation.
-
-        Returns:
-            None
+            licq_min (float): Minimum LICQ constraint value. Defaults to 0.0.
+            licq_max (float): Maximum LICQ constraint value. Defaults to 1e-4.
+            time_dilation_factor_min (float): Minimum time dilation factor.
+                Defaults to 0.3.
+            time_dilation_factor_max (float): Maximum time dilation factor.
+                Defaults to 3.0.
+            byof (ByofSpec, optional): Expert mode only. Raw JAX functions to
+                bypass symbolic layer. See :class:`openscvx.expert.ByofSpec` for
+                detailed documentation.
 
         Note:
             There are two approaches for handling time:
@@ -236,7 +236,7 @@ class Problem:
         self._algorithm = PenalizedTrustRegion()
 
     @property
-    def parameters(self):
+    def parameters(self) -> ParameterDict:
         """Get the parameters dictionary.
 
         The returned dictionary automatically syncs to CVXPy when modified:
@@ -244,7 +244,7 @@ class Problem:
             problem.parameters.update({"gate_0_center": center})  # Also syncs
 
         Returns:
-            ParameterDict: Special dict that syncs to CVXPy on assignment
+            ParameterDict: Special dict that syncs to CVXPy on assignment.
         """
         return self._parameter_wrapper
 
@@ -314,11 +314,13 @@ class Problem:
         index into the unified x and u vectors.
 
         Returns:
-            dict[str, slice]: Dictionary mapping variable names to slice objects.
+            Dictionary mapping variable names to slice objects.
                 State variables map to slices in the x vector.
                 Control variables map to slices in the u vector.
 
         Example:
+            Usage with byof::
+
                 problem = ox.Problem(dynamics, states, controls, ...)
                 print(problem.slices)
                 # {'position': slice(0, 3), 'velocity': slice(3, 6), 'theta': slice(0, 1)}

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/propagation/propagation.py

@@ -16,7 +16,7 @@ def prop_aug_dy(
     state_dot: callable,
     dis_type: str,
     N: int,
-    params,
+    params: dict,
 ) -> np.ndarray:
     """Compute the augmented dynamics for propagation.
 
@@ -50,16 +50,15 @@ def prop_aug_dy(
     return u[:, idx_s] * state_dot(x, u[:, :-1], node, params).squeeze()
 
 
-def get_propagation_solver(state_dot: Dynamics, settings: Config):
+def get_propagation_solver(state_dot: Dynamics, settings: Config) -> callable:
     """Create a propagation solver function.
 
     This function creates a solver that propagates the system state using the
     specified dynamics and settings.
 
     Args:
-        state_dot (callable): Function computing state derivatives.
+        state_dot: Dynamics object containing state derivative function.
         settings: Configuration settings for propagation.
-        param_map (dict): Mapping of parameter names to values.
 
     Returns:
         callable: A function that solves the propagation problem.
@@ -97,7 +96,7 @@ def get_propagation_solver(state_dot: Dynamics, settings: Config):
     return propagation_solver
 
 
-def s_to_t(x: np.ndarray, u: np.ndarray, settings: Config):
+def s_to_t(x: np.ndarray, u: np.ndarray, settings: Config) -> list[float]:
     """Convert normalized time s to real time t.
 
     This function converts the normalized time variable s to real time t
@@ -109,7 +108,7 @@ def s_to_t(x: np.ndarray, u: np.ndarray, settings: Config):
         settings (Config): Configuration settings.
 
     Returns:
-        list: List of real time points.
+        list[float]: List of real time points.
     """
     t = [x[:, settings.sim.time_slice][0]]
     tau = np.linspace(0, 1, settings.scp.n)
@@ -123,7 +122,9 @@ def s_to_t(x: np.ndarray, u: np.ndarray, settings: Config):
     return t
 
 
-def t_to_tau(u: np.ndarray, t, t_nodal, settings: Config):
+def t_to_tau(
+    u: np.ndarray, t: np.ndarray, t_nodal: np.ndarray, settings: Config
+) -> tuple[np.ndarray, np.ndarray]:
     """Convert real time t to normalized time tau.
 
     This function converts real time t to normalized time tau and interpolates
@@ -136,7 +137,8 @@ def t_to_tau(u: np.ndarray, t, t_nodal, settings: Config):
         settings (Config): Configuration settings.
 
     Returns:
-        tuple: (tau, u_interp) where tau is normalized time and u_interp is interpolated controls.
+        tuple[np.ndarray, np.ndarray]: (tau, u_interp) where tau is normalized time and u_interp is
+            interpolated controls.
     """
     if settings.dis.dis_type == "ZOH":
         # Zero-Order Hold: step interpolation (hold previous value)
@@ -170,7 +172,15 @@ def t_to_tau(u: np.ndarray, t, t_nodal, settings: Config):
     return tau, u_interp
 
 
-def simulate_nonlinear_time(params, x, u, tau_vals, t, settings, propagation_solver):
+def simulate_nonlinear_time(
+    params: dict,
+    x: np.ndarray,
+    u: np.ndarray,
+    tau_vals: np.ndarray,
+    t: np.ndarray,
+    settings: Config,
+    propagation_solver: callable,
+) -> np.ndarray:
     """Simulate the nonlinear system dynamics over time.
 
     This function simulates the system dynamics using the optimal control sequence

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/solvers/base.py

@@ -30,7 +30,7 @@ must follow for use within successive convexification algorithms.
 """
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, Any, List
 
 if TYPE_CHECKING:
     from openscvx.config import Config
@@ -197,7 +197,7 @@ class ConvexSolver(ABC):
         ...
 
     @abstractmethod
-    def solve(self):
+    def solve(self) -> Any:
         """Solve the convex subproblem and return results.
 
         Called at each SCP iteration after updating linearization and penalties.

{openscvx-0.3.2.dev316 → openscvx-0.3.2.dev317}/openscvx/symbolic/expr/arithmetic.py

@@ -32,7 +32,7 @@ Example:
         b = A @ x           # Creates MatMul(A, x)
 """
 
-from typing import Tuple
+from typing import Tuple, Union
 
 import numpy as np
 
@@ -56,7 +56,7 @@ class Add(Expr):
             z = x + y + 5  # Creates Add(x, y, Constant(5))
     """
 
-    def __init__(self, *args):
+    def __init__(self, *args: Union[Expr, float, int, np.ndarray]):
         """Initialize an addition operation.
 
         Args:
@@ -117,7 +117,7 @@ class Add(Expr):
         except ValueError as e:
             raise ValueError(f"Add shapes not broadcastable: {shapes}") from e
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         inner = " + ".join(repr(e) for e in self.terms)
         return f"({inner})"
 
@@ -140,7 +140,11 @@ class Sub(Expr):
             z = x - y  # Creates Sub(x, y)
     """
 
-    def __init__(self, left, right):
+    def __init__(
+        self,
+        left: Union[Expr, float, int, np.ndarray],
+        right: Union[Expr, float, int, np.ndarray],
+    ):
         """Initialize a subtraction operation.
 
         Args:
@@ -180,7 +184,7 @@ class Sub(Expr):
         except ValueError as e:
             raise ValueError(f"Sub shapes not broadcastable: {shapes}") from e
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"({self.left!r} - {self.right!r})"
 
 
@@ -202,7 +206,7 @@ class Mul(Expr):
             z = x * y * 2  # Creates Mul(x, y, Constant(2))
     """
 
-    def __init__(self, *args):
+    def __init__(self, *args: Union[Expr, float, int, np.ndarray]):
         """Initialize an element-wise multiplication operation.
 
         Args:
@@ -275,7 +279,7 @@ class Mul(Expr):
         except ValueError as e:
             raise ValueError(f"Mul shapes not broadcastable: {shapes}") from e
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         inner = " * ".join(repr(e) for e in self.factors)
         return f"({inner})"
 
@@ -298,7 +302,11 @@ class Div(Expr):
             z = x / y  # Creates Div(x, y)
     """
 
-    def __init__(self, left, right):
+    def __init__(
+        self,
+        left: Union[Expr, float, int, np.ndarray],
+        right: Union[Expr, float, int, np.ndarray],
+    ):
         """Initialize a division operation.
 
         Args:
@@ -338,7 +346,7 @@ class Div(Expr):
         except ValueError as e:
             raise ValueError(f"Div shapes not broadcastable: {shapes}") from e
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"({self.left!r} / {self.right!r})"
 
 
@@ -364,7 +372,11 @@ class MatMul(Expr):
             y = A @ x  # Creates MatMul(A, x), result shape (3,)
     """
 
-    def __init__(self, left, right):
+    def __init__(
+        self,
+        left: Union[Expr, float, int, np.ndarray],
+        right: Union[Expr, float, int, np.ndarray],
+    ):
         """Initialize a matrix multiplication operation.
 
         Args:
@@ -416,7 +428,7 @@ class MatMul(Expr):
                 raise ValueError(f"MatMul incompatible: {L} @ {R}")
             return L[:-1] + (R[-1],)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"({self.left!r} * {self.right!r})"
 
 
@@ -436,7 +448,7 @@ class Neg(Expr):
             y = -x  # Creates Neg(x)
     """
 
-    def __init__(self, operand):
+    def __init__(self, operand: Union[Expr, float, int, np.ndarray]):
         """Initialize a negation operation.
 
         Args:
@@ -462,7 +474,7 @@ class Neg(Expr):
         """Negation preserves the shape of its operand."""
         return self.operand.check_shape()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"(-{self.operand!r})"
 
 
@@ -483,7 +495,11 @@ class Power(Expr):
             y = x ** 2  # Creates Power(x, Constant(2))
     """
 
-    def __init__(self, base, exponent):
+    def __init__(
+        self,
+        base: Union[Expr, float, int, np.ndarray],
+        exponent: Union[Expr, float, int, np.ndarray],
+    ):
         """Initialize a power operation.
 
         Args:
@@ -513,5 +529,5 @@ class Power(Expr):
         except ValueError as e:
             raise ValueError(f"Power shapes not broadcastable: {shapes}") from e
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"({self.base!r})**({self.exponent!r})"