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

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

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

## Module Contents

### Functions

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

* - {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}`tree_device_put <evorl.distributed.sharding.tree_device_put>`
  - ```{autodoc2-docstring} evorl.distributed.sharding.tree_device_put
    :parser: autodoc2_docstrings_parser
    :summary:
    ```
````

### API

````{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} 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
```
````