Home
last modified time | relevance | path

Searched refs:numpy_nonzero (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcustom_op_db.py202 def numpy_nonzero(x: Tensor) -> Tensor: function
209 @numpy_nonzero.register_fake
230 numpy_nonzero.register_vmap(numpy_nonzero_vmap)
482 op=numpy_nonzero._opoverload,
/aosp_15_r20/external/pytorch/test/
H A Dtest_custom_ops.py40 from torch.testing._internal.custom_op_db import numpy_nonzero
1569 numpy_nonzero(x)
1611 gm = make_fx(numpy_nonzero, tracing_mode="symbolic")(x)
1618 make_fx(numpy_nonzero, tracing_mode="fake")(x)
1649 return numpy_nonzero(x.clone()).clone()