Grad¶

tinax.grad covers JAX's autodiff surface with hardened argnums validation shared across all three transforms: value_and_grad, jacobian (an explicit mode="forward" or "reverse" in place of separate jacfwd/jacrev calls), and hessian.

from tinax.grad import value_and_grad

loss_and_grad = value_and_grad(loss)
value, gradient = loss_and_grad(params, batch)

Raw jax.grad silently accepts a boolean argnums, differentiating with respect to the wrong argument with no error. tinax.grad rejects it, along with duplicate and empty selections.

value_and_grad ¶

value_and_grad(
    function: Callable[P, object],
    *,
    argnums: int | Sequence[int] = 0,
    has_aux: bool = False,
) -> Callable[P, tuple[object, object]]

Differentiate function, validating argnums where raw jax.grad silently accepts a bool.

Bare grad is not exposed: value_and_grad is a strict superset, and callers who only need the gradient can discard the value.

Parameters:

Name	Type	Description	Default
`function`	`Callable[P, object]`	Callable to differentiate.	required
`argnums`	`int \| Sequence[int]`	Positional argument index or indices to differentiate with respect to. Must be nonnegative integers with no duplicates. Defaults to 0.	`0`
`has_aux`	`bool`	Whether `function` returns `(output, auxiliary)`. Defaults to `False`.	`False`

Returns:

Type	Description
`Callable[P, tuple[object, object]]`	A callable returning `(value, gradient)` for the selected arguments.

Raises:

Type	Description
`TypeError`	If `function` is not callable, `argnums` is not an integer or sequence of integers (booleans rejected), or `has_aux` is not a bool.
`ValueError`	If `argnums` is negative, empty (as a sequence), or contains a duplicate index.

jacobian ¶

jacobian(
    function: Callable[P, object],
    *,
    mode: Literal["forward", "reverse"],
    argnums: int | Sequence[int] = 0,
    has_aux: bool = False,
) -> Callable[P, object]

Compute a Jacobian, requiring an explicit mode instead of a separate jacfwd/jacrev call.

Raw JAX splits this into two separately named top-level functions with a real wide-versus-tall performance tradeoff; mode turns that implicit choice into one explicit, validated parameter.

Parameters:

Name	Type	Description	Default
`function`	`Callable[P, object]`	Callable to differentiate.	required
`mode`	`Literal['forward', 'reverse']`	`"forward"` uses forward-mode autodiff (`jax.jacfwd`), efficient for few inputs and many outputs. `"reverse"` uses reverse-mode autodiff (`jax.jacrev`), efficient for many inputs and few outputs.	required
`argnums`	`int \| Sequence[int]`	Positional argument index or indices to differentiate with respect to. Must be nonnegative integers with no duplicates. Defaults to 0.	`0`
`has_aux`	`bool`	Whether `function` returns `(output, auxiliary)`. Defaults to `False`.	`False`

Returns:

Type	Description
`Callable[P, object]`	A callable returning the Jacobian for the selected arguments.

Raises:

Type	Description
`TypeError`	If `function` is not callable, `mode` is not a string, `argnums` is not an integer or sequence of integers (booleans rejected), or `has_aux` is not a bool.
`ValueError`	If `mode` is not `"forward"` or `"reverse"`, `argnums` is negative, empty (as a sequence), or contains a duplicate index.

hessian ¶

hessian(
    function: Callable[P, object],
    *,
    argnums: int | Sequence[int] = 0,
) -> Callable[P, object]

Compute a Hessian, sharing value_and_grad's and jacobian's argnums validation.

Parameters:

Name	Type	Description	Default
`function`	`Callable[P, object]`	Callable to differentiate twice.	required
`argnums`	`int \| Sequence[int]`	Positional argument index or indices to differentiate with respect to. Must be nonnegative integers with no duplicates. Defaults to 0.	`0`

Returns:

Type	Description
`Callable[P, object]`	A callable returning the Hessian for the selected arguments.

Raises:

Type	Description
`TypeError`	If `function` is not callable, or `argnums` is not an integer or sequence of integers (booleans rejected).
`ValueError`	If `argnums` is negative, empty (as a sequence), or contains a duplicate index.