diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py index 78018a3e080..ea8f08a86fa 100644 --- a/tools/pyi/gen_pyi.py +++ b/tools/pyi/gen_pyi.py @@ -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", + ) + ], } )