OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Add api info for torch._C._nn.pyi (#148405)

Browse Source 
APis involved are as followed:

- adaptive_avg_pool2d
- adaptive_avg_pool3d
- binary_cross_entropy
- col2im

ISSUE Related:
https://github.com/pytorch/pytorch/issues/148404
Pull Request resolved: https://github.com/pytorch/pytorch/pull/148405
Approved by: https://github.com/ezyang
This commit is contained in:
FFFrog 2025-06-13 18:10:40 +08:00 committed by PyTorch MergeBot
parent 7070ab3180
commit 1982ec2d22
1 changed files with 35 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
tools/pyi/gen_pyi.py
View File

@ -412,6 +412,16 @@ def gen_nn_functional(fm: FileManager) -> None:
"tuple[Tensor, Tensor]",
)
],
f"adaptive_avg_pool{d}d": [
defs(
f"adaptive_avg_pool{d}d",
[
INPUT,
"output_size: _int | _size",
],
"Tensor",
)
],
}
)
@ -516,6 +526,31 @@ def gen_nn_functional(fm: FileManager) -> None:
"Tensor",
)
],
"binary_cross_entropy": [
defs(
"binary_cross_entropy",
[
INPUT,
"target: Tensor",
"weight: Tensor | None = None",
"reduction: str = ...",
],
"Tensor",
)
],
"col2im": [
defs(
"col2im",
[
INPUT,
"output_size: _int | _size",
KERNEL_SIZE,
"dilation: _int | _size",
*STRIDE_PADDING,
],
"Tensor",
)
],
}
)