xref: /aosp_15_r20/external/pytorch/functorch/_src/vmap/__init__.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# This file has moved to under torch/_functorch. It is not public API.
2# If you are not a PyTorch developer and you are relying on the following
3# imports, please file an issue.
4from torch._functorch.vmap import (
5    _add_batch_dim,
6    _broadcast_to_and_flatten,
7    _create_batched_inputs,
8    _get_name,
9    _process_batched_inputs,
10    _remove_batch_dim,
11    _unwrap_batched,
12    _validate_and_get_batch_size,
13    Tensor,
14    tree_flatten,
15    tree_unflatten,
16)
17