# {py:mod}`evorl.distributed.gradients`

```{py:module} evorl.distributed.gradients
```

```{autodoc2-docstring} evorl.distributed.gradients
:parser: autodoc2_docstrings_parser
:allowtitles:
```

## Module Contents

### Functions

````{list-table}
:class: autosummary longtable
:align: left

* - {py:obj}`agent_gradient_update <evorl.distributed.gradients.agent_gradient_update>`
  - ```{autodoc2-docstring} evorl.distributed.gradients.agent_gradient_update
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`gradient_update <evorl.distributed.gradients.gradient_update>`
  - ```{autodoc2-docstring} evorl.distributed.gradients.gradient_update
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`loss_and_pgrad <evorl.distributed.gradients.loss_and_pgrad>`
  - ```{autodoc2-docstring} evorl.distributed.gradients.loss_and_pgrad
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

````{py:function} agent_gradient_update(loss_fn: collections.abc.Callable[..., float], optimizer: optax.GradientTransformation, dp_axis_name: str | None = None, has_aux: bool = False, attach_fn: collections.abc.Callable[[chex.ArrayTree, chex.ArrayTree], chex.ArrayTree] = _attach_params_to_agent_state, detach_fn: collections.abc.Callable[[chex.ArrayTree, chex.ArrayTree], chex.ArrayTree] = _detach_params_to_agent_state)
:canonical: evorl.distributed.gradients.agent_gradient_update

```{autodoc2-docstring} evorl.distributed.gradients.agent_gradient_update
:parser: autodoc2_docstrings_parser
```
````

````{py:function} gradient_update(loss_fn: collections.abc.Callable[..., float], optimizer: optax.GradientTransformation, dp_axis_name: str | None, has_aux: bool = False)
:canonical: evorl.distributed.gradients.gradient_update

```{autodoc2-docstring} evorl.distributed.gradients.gradient_update
:parser: autodoc2_docstrings_parser
```
````

````{py:function} loss_and_pgrad(loss_fn: collections.abc.Callable[..., float], dp_axis_name: str | None, has_aux: bool = False)
:canonical: evorl.distributed.gradients.loss_and_pgrad

```{autodoc2-docstring} evorl.distributed.gradients.loss_and_pgrad
:parser: autodoc2_docstrings_parser
```
````