# {py:mod}`evorl.distributed`

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

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

## Package 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}`all_gather <evorl.distributed.comm.all_gather>`
  - ```{autodoc2-docstring} evorl.distributed.comm.all_gather
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`get_global_ranks <evorl.distributed.comm.get_global_ranks>`
  - ```{autodoc2-docstring} evorl.distributed.comm.get_global_ranks
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`get_process_id <evorl.distributed.comm.get_process_id>`
  - ```{autodoc2-docstring} evorl.distributed.comm.get_process_id
    :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}`is_dist_initialized <evorl.distributed.comm.is_dist_initialized>`
  - ```{autodoc2-docstring} evorl.distributed.comm.is_dist_initialized
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`pmax <evorl.distributed.comm.pmax>`
  - ```{autodoc2-docstring} evorl.distributed.comm.pmax
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`pmean <evorl.distributed.comm.pmean>`
  - ```{autodoc2-docstring} evorl.distributed.comm.pmean
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`pmin <evorl.distributed.comm.pmin>`
  - ```{autodoc2-docstring} evorl.distributed.comm.pmin
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`psum <evorl.distributed.comm.psum>`
  - ```{autodoc2-docstring} evorl.distributed.comm.psum
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`shmap_map <evorl.distributed.sharding.shmap_map>`
  - ```{autodoc2-docstring} evorl.distributed.sharding.shmap_map
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`shmap_vmap <evorl.distributed.sharding.shmap_vmap>`
  - ```{autodoc2-docstring} evorl.distributed.sharding.shmap_vmap
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`split_key_to_devices <evorl.distributed.comm.split_key_to_devices>`
  - ```{autodoc2-docstring} evorl.distributed.comm.split_key_to_devices
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`tree_device_put <evorl.distributed.sharding.tree_device_put>`
  - ```{autodoc2-docstring} evorl.distributed.sharding.tree_device_put
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### Data

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

* - {py:obj}`DP_AXIS_NAME <evorl.distributed.DP_AXIS_NAME>`
  - ```{autodoc2-docstring} evorl.distributed.DP_AXIS_NAME
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
* - {py:obj}`POP_AXIS_NAME <evorl.distributed.POP_AXIS_NAME>`
  - ```{autodoc2-docstring} evorl.distributed.POP_AXIS_NAME
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

````{py:data} DP_AXIS_NAME
:canonical: evorl.distributed.DP_AXIS_NAME
:value: >
   'DP'

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

````

````{py:data} POP_AXIS_NAME
:canonical: evorl.distributed.POP_AXIS_NAME
:value: >
   'POP'

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

````

````{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} all_gather(x, axis_name: str | None = None, **kwargs)
:canonical: evorl.distributed.comm.all_gather

```{autodoc2-docstring} evorl.distributed.comm.all_gather
:parser: autodoc2_docstrings_parser
```
````

````{py:function} get_global_ranks()
:canonical: evorl.distributed.comm.get_global_ranks

```{autodoc2-docstring} evorl.distributed.comm.get_global_ranks
:parser: autodoc2_docstrings_parser
```
````

````{py:function} get_process_id()
:canonical: evorl.distributed.comm.get_process_id

```{autodoc2-docstring} evorl.distributed.comm.get_process_id
: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} is_dist_initialized()
:canonical: evorl.distributed.comm.is_dist_initialized

```{autodoc2-docstring} evorl.distributed.comm.is_dist_initialized
:parser: autodoc2_docstrings_parser
```
````

````{py:function} pmax(x, axis_name: str | None = None)
:canonical: evorl.distributed.comm.pmax

```{autodoc2-docstring} evorl.distributed.comm.pmax
:parser: autodoc2_docstrings_parser
```
````

````{py:function} pmean(x, axis_name: str | None = None)
:canonical: evorl.distributed.comm.pmean

```{autodoc2-docstring} evorl.distributed.comm.pmean
:parser: autodoc2_docstrings_parser
```
````

````{py:function} pmin(x, axis_name: str | None = None)
:canonical: evorl.distributed.comm.pmin

```{autodoc2-docstring} evorl.distributed.comm.pmin
:parser: autodoc2_docstrings_parser
```
````

````{py:function} psum(x, axis_name: str | None = None)
:canonical: evorl.distributed.comm.psum

```{autodoc2-docstring} evorl.distributed.comm.psum
:parser: autodoc2_docstrings_parser
```
````

````{py:function} shmap_map(fn: collections.abc.Callable, mesh, in_specs, out_specs, **kwargs)
:canonical: evorl.distributed.sharding.shmap_map

```{autodoc2-docstring} evorl.distributed.sharding.shmap_map
:parser: autodoc2_docstrings_parser
```
````

````{py:function} shmap_vmap(fn: collections.abc.Callable, mesh, in_specs, out_specs, **kwargs)
:canonical: evorl.distributed.sharding.shmap_vmap

```{autodoc2-docstring} evorl.distributed.sharding.shmap_vmap
:parser: autodoc2_docstrings_parser
```
````

````{py:function} split_key_to_devices(key: chex.PRNGKey, devices: collections.abc.Sequence[jax.Device])
:canonical: evorl.distributed.comm.split_key_to_devices

```{autodoc2-docstring} evorl.distributed.comm.split_key_to_devices
:parser: autodoc2_docstrings_parser
```
````

````{py:function} tree_device_put(tree: chex.ArrayTree, device_or_sharding)
:canonical: evorl.distributed.sharding.tree_device_put

```{autodoc2-docstring} evorl.distributed.sharding.tree_device_put
:parser: autodoc2_docstrings_parser
```
````