Skip to content

Commit

Permalink
add docstring to root_scalar
Browse files Browse the repository at this point in the history
  • Loading branch information
Joseph Kleinhenz committed Aug 28, 2024
1 parent 4764b34 commit fc3fd0f
Showing 1 changed file with 28 additions and 2 deletions.
30 changes: 28 additions & 2 deletions src/beignet/_root_scalar.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,36 @@ class RootSolutionInfo:
def root_scalar(
func: Callable,
*args,
method: Literal["bisect"] | Literal["chandrupatla"] = "chandrupatla",
method: Literal["bisect", "chandrupatla"] = "chandrupatla",
implicit_diff: bool = True,
options: dict,
options: dict | None = None,
):
"""
Find the root of a scalar (elementwise) function.
Parameters
----------
func: Callable
Function to find a root of. Called as `f(x, *args)`.
The function must operate element wise, i.e. `f(x[i]) == f(x)[i]`.
Handling *args via broadcasting is acceptable.
*args
Extra arguments to be passed to `func`.
method: Literal["bisect", "chandrupatla"] = "chandrupatla"
Solver method to use.
implicit_diff: bool = True
If true, the solver is wrapped in `beignet.func.custom_scalar_root` which
enables gradients with respect to *args using implicit differentiation.
options: dict | None = None
A dictionary of options that are passed through to the solver as keyword args.
"""
if options is None:
options = {}

if method == "bisect":
solver = beignet.bisect
elif method == "chandrupatla":
Expand Down

0 comments on commit fc3fd0f

Please sign in to comment.