Searched refs:numpy_sort (Results 1 – 2 of 2) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/testing/_internal/ |
| H A D | custom_op_db.py | 117 def numpy_sort(x: Tensor, dim: int) -> Tuple[Tensor, Tensor, Tensor]: function
129 @numpy_sort.register_fake
143 numpy_sort.register_autograd(numpy_sort_backward, setup_context=numpy_sort_setup_context)
149 result = numpy_sort(x, dim + 1)
152 numpy_sort.register_vmap(numpy_sort_vmap)
468 op=numpy_sort._opoverload,
|
| /aosp_15_r20/external/pytorch/docs/source/notes/ |
| H A D | extending.func.rst | 158 def numpy_sort(x, dim=-1):
165 grad_x = torch.func.grad(lambda x: numpy_sort(x).sum())(x)
475 def numpy_sort(x, dim=-1):
480 result = torch.vmap(numpy_sort)(x)
481 assert torch.allclose(result, numpy_sort(result, 1))
|