pytorch/functorch/_src/vmap/__init__.py
Edward Z. Yang e6ec0efaf8 Apply UFMT to all non test/torch files (#106205)
Signed-off-by: Edward Z. Yang <ezyang@meta.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/106205
Approved by: https://github.com/albanD
2023-07-29 02:56:24 +00:00

16 lines
467 B
Python

# This file has moved to under torch/_functorch. It is not public API.
# If you are not a PyTorch developer and you are relying on the following
# imports, please file an issue.
from torch._functorch.vmap import (
_add_batch_dim,
_broadcast_to_and_flatten,
_create_batched_inputs,
_get_name,
_process_batched_inputs,
_remove_batch_dim,
_unwrap_batched,
_validate_and_get_batch_size,
Tensor,
tree_flatten,
tree_unflatten,
)