jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _compare_onnx_and_torch_exported_program(
        self,
        torch_exported_program,
        onnx_exported_program,
        input_args,
        input_kwargs=None,
        rtol=1e-03,
        atol=1e-07,
    ):
        # avoid mutable default argument
        if input_kwargs is None:
            input_kwargs = {}

        # NOTE: ONNXProgram holds a reference (not copy) to the original ref_model, including its state_dict.
        # Thus, ONNXProgram() must run before ref_model() to prevent ref_model.forward() from changing the state_dict.
        # Otherwise, the ref_model can change buffers on state_dict which would be used by ONNXProgram.__call__()
        onnx_outputs = onnx_exported_program(*input_args, **input_kwargs)
        if isinstance(torch_exported_program, torch.export.ExportedProgram):
            torch_outputs = torch_exported_program.module()(*input_args, **input_kwargs)
        else:
            torch_outputs = torch_exported_program(*input_args, **input_kwargs)
        torch_outputs_onnx_format = onnx_exported_program.adapt_torch_outputs_to_onnx(
            torch_outputs
        )
        if len(torch_outputs_onnx_format) != len(onnx_outputs):
            raise AssertionError(
                f"Expected {len(torch_outputs_onnx_format)} outputs, got {len(onnx_outputs)}"
            )
        for torch_output, onnx_output in zip(torch_outputs_onnx_format, onnx_outputs):
            torch.testing.assert_allclose(
                torch_output, torch.tensor(onnx_output), rtol=rtol, atol=atol
            )
//  [python]
def _compare_onnx_and_torch_exported_program(
        self,
        torch_exported_program,
        onnx_exported_program,
        input_args,
        input_kwargs=None,
        rtol=1e-03,
        atol=1e-07,
    ):
        # avoid mutable default argument
        if input_kwargs is None:
            input_kwargs = {}

        # NOTE: ONNXProgram holds a reference (not copy) to the original ref_model, including its state_dict.
        # Thus, ONNXProgram() must run before ref_model() to prevent ref_model.forward() from changing the state_dict.
        # Otherwise, the ref_model can change buffers on state_dict which would be used by ONNXProgram.__call__()
        onnx_outputs = onnx_exported_program(*input_args, **input_kwargs)
        if isinstance(torch_exported_program, torch.export.ExportedProgram):
            torch_outputs = torch_exported_program.module()(*input_args, **input_kwargs)
        else:
            torch_outputs = torch_exported_program(*input_args, **input_kwargs)
        torch_outputs_onnx_format = onnx_exported_program.adapt_torch_outputs_to_onnx(
            torch_outputs
        )
        if len(torch_outputs_onnx_format) != len(onnx_outputs):
            raise AssertionError(
                f"Expected {len(torch_outputs_onnx_format)} outputs, got {len(onnx_outputs)}"
            )
        for torch_output, onnx_output in zip(torch_outputs_onnx_format, onnx_outputs):
            torch.testing.assert_close(
                torch_output, torch.tensor(onnx_output), rtol=rtol, atol=atol
            )
//  [python]
def test_scan2d_for():

    @triton.jit
    def fn(out_ptr0, rnumel, RBLOCK: tl.constexpr):
        rbase = tl.arange(0, RBLOCK)[None, :]
        for roffset in range(0, rnumel, RBLOCK):
            rindex = roffset + rbase
            rmask = rindex < rnumel
            tmp3 = tl.where(rmask, 1, 0)
            tmp6 = tl.cumsum(tmp3, 1)
            tl.store(out_ptr0 + rindex, tmp6, rmask)

    RBLOCK = 8
    out0 = torch.empty(RBLOCK, device="cuda", dtype=torch.int64)
    fn[(1, )](out0, RBLOCK, RBLOCK)
    ref = torch.arange(RBLOCK, device="cuda", dtype=torch.int64) + 1
    torch.testing.assert_allclose(out0, ref)
//  [python]
def test_scan2d_for():

    @triton.jit
    def fn(out_ptr0, rnumel, RBLOCK: tl.constexpr):
        rbase = tl.arange(0, RBLOCK)[None, :]
        for roffset in range(0, rnumel, RBLOCK):
            rindex = roffset + rbase
            rmask = rindex < rnumel
            tmp3 = tl.where(rmask, 1, 0)
            tmp6 = tl.cumsum(tmp3, 1)
            tl.store(out_ptr0 + rindex, tmp6, rmask)

    RBLOCK = 8
    out0 = torch.empty(RBLOCK, device="cuda", dtype=torch.int64)
    fn[(1, )](out0, RBLOCK, RBLOCK)
    ref = torch.arange(RBLOCK, device="cuda", dtype=torch.int64) + 1
    torch.testing.assert_close(out0, ref)
//  [python]
def register_quantization_mapping_with_decorator_example():
    # ARRANGE
    class MyBlock(nn.Module):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.in_feats = in_feats
            self.out_feats = out_feats
            self.flatten = nn.Flatten()
            self.linear = nn.Linear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    @register_quantized_module(float_source=MyBlock)
    class MyQuantizedBlock(SGQuantMixin):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.flatten = nn.Flatten()
            self.linear = quant_nn.QuantLinear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    class MyModel(nn.Module):
        def __init__(self, res, n_classes) -> None:
            super().__init__()
            self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
            self.my_block = MyBlock(4 * (res**2), n_classes)

        def forward(self, x):
            y = self.conv(x)
            return self.my_block(y)

    res = 32
    n_clss = 10
    module = MyModel(res, n_clss)

    # TEST
    q_util = SelectiveQuantizer()
    q_util.quantize_module(module)

    x = torch.rand(1, 3, res, res)

    print(module)

    # ASSERT
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_allclose(y.size(), (1, n_clss))
//  [python]
def register_quantization_mapping_with_decorator_example():
    # ARRANGE
    class MyBlock(nn.Module):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.in_feats = in_feats
            self.out_feats = out_feats
            self.flatten = nn.Flatten()
            self.linear = nn.Linear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    @register_quantized_module(float_source=MyBlock)
    class MyQuantizedBlock(SGQuantMixin):
        def __init__(self, in_feats, out_feats) -> None:
            super().__init__()
            self.flatten = nn.Flatten()
            self.linear = quant_nn.QuantLinear(in_feats, out_feats)

        def forward(self, x):
            return self.linear(self.flatten(x))

    class MyModel(nn.Module):
        def __init__(self, res, n_classes) -> None:
            super().__init__()
            self.conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
            self.my_block = MyBlock(4 * (res**2), n_classes)

        def forward(self, x):
            y = self.conv(x)
            return self.my_block(y)

    res = 32
    n_clss = 10
    module = MyModel(res, n_clss)

    # TEST
    q_util = SelectiveQuantizer()
    q_util.quantize_module(module)

    x = torch.rand(1, 3, res, res)

    print(module)

    # ASSERT
    with torch.no_grad():
        y = module(x)
        torch.testing.assert_close(y.size(), (1, n_clss))
//  [python]
def test_apex_cross_entropy(device: str, dtype: torch.dtype):
    logits = torch.randn([2048, 50257], device=device, dtype=thunder.torch.to_torch_dtype(dtype))
    labels = torch.randint(0, 50257, [2048], device=device)

    def fn(logits, labels):
        return torch.nn.functional.cross_entropy(logits, labels, reduction="mean", ignore_index=-1)

    cfn = thunder.jit(fn, executors=[apex_ex])

    # Verifies the result is close to PyTorch
    thunder_result = cfn(logits, labels)
    torch_result = fn(logits, labels)

    t.assert_close(thunder_result, torch_result)

    # Verifies apex_cross_entropy was called
    extrace = thunder.last_traces(cfn)[-1]
    assert any(bsym.sym.name == "apex_cross_entropy" for bsym in extrace.bound_symbols)
//  [python]
def test_apex_cross_entropy(device: str, dtype: torch.dtype):
    logits = torch.randn([2048, 50257], device=device, dtype=thunder.torch.to_torch_dtype(dtype))
    labels = torch.randint(0, 50257, [2048], device=device)

    def fn(logits, labels):
        return torch.nn.functional.cross_entropy(logits, labels, reduction="mean", ignore_index=-1)

    cfn = thunder.jit(fn, executors=[apex_ex])

    # Verifies the result is close to PyTorch
    thunder_result = cfn(logits, labels)
    torch_result = fn(logits, labels)

    assert_close(thunder_result, torch_result)

    # Verifies apex_cross_entropy was called
    extrace = thunder.last_traces(cfn)[-1]
    assert any(bsym.sym.name == "apex_cross_entropy" for bsym in extrace.bound_symbols)
//  [python]
def test_rrr_no_tf32(self):
        # Test accuracy with tf32 disabled
        # this test uses a smaller numerical tolerance level
        # than the others
        allow_tf32_bak = torch.backends.cuda.matmul.allow_tf32
        torch.backends.cuda.matmul.allow_tf32 = False
        try:
            test_dtype = torch.float32
            test_dtype_str = "float32"
            A = torch.rand((64, 64), dtype=test_dtype).cuda()
            B = torch.rand((64, 64), dtype=test_dtype).cuda()
            result_cuda = torch.matmul(A, B)

            target = detect_target(no_tf32=True)  # Disable tf32 for accuracy
            A_ait = Tensor(
                shape=[64, 64], dtype=test_dtype_str, name="input_0", is_input=True
            )
            B_ait = Tensor(
                shape=[64, 64], dtype=test_dtype_str, name="input_1", is_input=True
            )
            OP = ops.gemm_rrr()
            Y = OP(A_ait, B_ait)
            Y._attrs["name"] = "output_0"
            Y._attrs["is_output"] = True
            module = compile_model(Y, target, "./tmp", "gemm_rrr_no_tf32")
            inputs = {
                "input_0": A.clone().detach().cuda(),
                "input_1": B.clone().detach().cuda(),
            }
            result_ait = torch.empty([64, 64], dtype=test_dtype, device="cuda")
            module.run_with_tensors(inputs, [result_ait])
            torch.testing.assert_allclose(result_cuda, result_ait)
        finally:
            torch.backends.cuda.matmul.allow_tf32 = allow_tf32_bak
//  [python]
def test_rrr_no_tf32(self):
        # Test accuracy with tf32 disabled
        # this test uses a smaller numerical tolerance level
        # than the others
        allow_tf32_bak = torch.backends.cuda.matmul.allow_tf32
        torch.backends.cuda.matmul.allow_tf32 = False
        try:
            test_dtype = torch.float32
            test_dtype_str = "float32"
            A = torch.rand((64, 64), dtype=test_dtype).cuda()
            B = torch.rand((64, 64), dtype=test_dtype).cuda()
            result_cuda = torch.matmul(A, B)

            target = detect_target(no_tf32=True)  # Disable tf32 for accuracy
            A_ait = Tensor(
                shape=[64, 64], dtype=test_dtype_str, name="input_0", is_input=True
            )
            B_ait = Tensor(
                shape=[64, 64], dtype=test_dtype_str, name="input_1", is_input=True
            )
            OP = ops.gemm_rrr()
            Y = OP(A_ait, B_ait)
            Y._attrs["name"] = "output_0"
            Y._attrs["is_output"] = True
            module = compile_model(Y, target, "./tmp", "gemm_rrr_no_tf32")
            inputs = {
                "input_0": A.clone().detach().cuda(),
                "input_1": B.clone().detach().cuda(),
            }
            result_ait = torch.empty([64, 64], dtype=test_dtype, device="cuda")
            module.run_with_tensors(inputs, [result_ait])
            torch.testing.assert_close(result_cuda, result_ait)
        finally:
            torch.backends.cuda.matmul.allow_tf32 = allow_tf32_bak
//  [python]
def _test_quantize_dequantize(
        self,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 3.5,
        dtype: torch.dtype = torch.float32,
        dims: DimsType = 23,
    ) -> None:
        """Check numerical error when casting to FP8 and back"""

        # Initialize random data
        x_ref = 2 * torch.rand(_to_list(dims), dtype=dtype, device="cpu") - 1

        # Cast to FP8 and back
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_fp8 = x_fp8.from_float8().cpu()

        # Check results
        t.assert_close(x_fp8, x_ref, **_tols[fp8_dtype])

        # Make sure we are not trivially passing the test
        with pytest.raises(AssertionError):
            t.assert_close(x_fp8, -x_ref, **_tols[fp8_dtype])
//  [python]
def _test_quantize_dequantize(
        self,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 3.5,
        dtype: torch.dtype = torch.float32,
        dims: DimsType = 23,
    ) -> None:
        """Check numerical error when casting to FP8 and back"""

        # Initialize random data
        x_ref = 2 * torch.rand(_to_list(dims), dtype=dtype, device="cpu") - 1

        # Cast to FP8 and back
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_fp8 = x_fp8.from_float8().cpu()

        # Check results
        torch.testing.assert_close(x_fp8, x_ref, **_tols[fp8_dtype])

        # Make sure we are not trivially passing the test
        with pytest.raises(AssertionError):
            torch.testing.assert_close(x_fp8, -x_ref, **_tols[fp8_dtype])
//  [python]
def test_apex_cross_entropy_backward(device, dtype):
    from thunder.core.transforms import value_and_grad

    logits = torch.randn([2048, 50257], device=device, dtype=thunder.torch.to_torch_dtype(dtype), requires_grad=True)
    labels = torch.randint(0, 50257, [2048], device=device)

    # -1 is supported by apex cross entropy but 1 is not. The case of 1 is
    # used to test that the conditional rules are working correctly and that
    # the apex cross entropy is not used
    ignore_indices = (-1, 1)

    for ignore_index in ignore_indices:

        @value_and_grad
        def test(logits, labels):
            return thunder.torch.cross_entropy(logits, labels, reduction="mean", ignore_index=ignore_index)

        ctest = thunder.compile(
            test,
            executors_list=[apex_ex],
            disable_preprocessing=True,
            disable_torch_autograd_support=True,
        )
        actual = ctest(logits, labels)
        expected = torch.nn.functional.cross_entropy(logits, labels, reduction="mean", ignore_index=ignore_index)
        expected_grad = torch.autograd.grad(expected, logits)[0]
        t.assert_close(actual[0], expected)
        t.assert_close(actual[1][0], expected_grad)
        last_trace = thunder.last_traces(ctest)[-1]
        is_any_fw = any(bsym.sym.name == "apex_cross_entropy" for bsym in last_trace.bound_symbols)
        is_any_bw = any(bsym.sym.name == "apex_cross_entropy_backward" for bsym in last_trace.bound_symbols)

        if ignore_index == -1:
            assert is_any_fw
            assert is_any_bw
        else:
            assert not is_any_fw
            assert not is_any_bw
//  [python]
def test_apex_cross_entropy_backward(device, dtype):
    from thunder.core.transforms import value_and_grad

    logits = torch.randn([2048, 50257], device=device, dtype=thunder.torch.to_torch_dtype(dtype), requires_grad=True)
    labels = torch.randint(0, 50257, [2048], device=device)

    # -1 is supported by apex cross entropy but 1 is not. The case of 1 is
    # used to test that the conditional rules are working correctly and that
    # the apex cross entropy is not used
    ignore_indices = (-1, 1)

    for ignore_index in ignore_indices:

        @value_and_grad
        def test(logits, labels):
            return thunder.torch.cross_entropy(logits, labels, reduction="mean", ignore_index=ignore_index)

        ctest = thunder.compile(
            test,
            executors_list=[apex_ex],
            disable_preprocessing=True,
            disable_torch_autograd_support=True,
        )
        actual = ctest(logits, labels)
        expected = torch.nn.functional.cross_entropy(logits, labels, reduction="mean", ignore_index=ignore_index)
        expected_grad = torch.autograd.grad(expected, logits)[0]
        torch.testing.assert_close(actual[0], expected)
        torch.testing.assert_close(actual[1][0], expected_grad)
        last_trace = thunder.last_traces(ctest)[-1]
        is_any_fw = any(bsym.sym.name == "apex_cross_entropy" for bsym in last_trace.bound_symbols)
        is_any_bw = any(bsym.sym.name == "apex_cross_entropy_backward" for bsym in last_trace.bound_symbols)

        if ignore_index == -1:
            assert is_any_fw
            assert is_any_bw
        else:
            assert not is_any_fw
            assert not is_any_bw
//  [python]
def torch_assert_equal(actual, expected, **kwargs):
    # assert_close was added around pt-1.9, it does better checks - e.g will check dimensions match
    if hasattr(torch.testing, "assert_close"):
        return torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0, **kwargs)
    else:
        return torch.allclose(actual, expected, rtol=0.0, atol=0.0)
//  [python]
def test_inplace_ops(
        self,
        dims: DimsType = 23,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 3.5,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        """Test in-place ops"""

        # Initialize random data
        dims = _to_list(dims)
        x_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        y_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        y_fp8 = Float8Tensor.to_float8(
            y_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_ref = x_fp8.from_float8()
        y_ref = y_fp8.from_float8()

        # In-place operations
        tols = _tols[fp8_dtype]
        x_fp8 += y_ref
        x_ref += y_ref
        t.assert_close(x_fp8, x_ref, **tols)
        x_ref = x_fp8.from_float8()
        x_fp8 -= y_fp8
        x_ref -= y_fp8
        t.assert_close(x_fp8, x_ref, **tols)
        x_ref = x_fp8.from_float8()
        x_fp8 *= 2
        x_ref *= 2
        t.assert_close(x_fp8, x_ref, **tols)
        x_ref = x_fp8.from_float8()

        # Make sure we are not trivially passing tests
        x_ref += 123
        with pytest.raises(AssertionError):
            t.assert_close(x_fp8, x_ref, **tols)
//  [python]
def test_inplace_ops(
        self,
        dims: DimsType = 23,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 3.5,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        """Test in-place ops"""

        # Initialize random data
        dims = _to_list(dims)
        x_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        y_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        y_fp8 = Float8Tensor.to_float8(
            y_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_ref = x_fp8.from_float8()
        y_ref = y_fp8.from_float8()

        # In-place operations
        tols = _tols[fp8_dtype]
        x_fp8 += y_ref
        x_ref += y_ref
        torch.testing.assert_close(x_fp8, x_ref, **tols)
        x_ref = x_fp8.from_float8()
        x_fp8 -= y_fp8
        x_ref -= y_fp8
        torch.testing.assert_close(x_fp8, x_ref, **tols)
        x_ref = x_fp8.from_float8()
        x_fp8 *= 2
        x_ref *= 2
        torch.testing.assert_close(x_fp8, x_ref, **tols)
        x_ref = x_fp8.from_float8()

        # Make sure we are not trivially passing tests
        x_ref += 123
        with pytest.raises(AssertionError):
            torch.testing.assert_close(x_fp8, x_ref, **tols)
//  [python]
def test_serialization(
        self,
        dims: DimsType = [2,3,5],
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 0.5,
        dtype: torch.dtype = torch.float32,
    ):

        # Initialize random data
        dims = _to_list(dims)
        x_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_ref = x_fp8.from_float8()

        # Serialize tensor
        byte_stream = io.BytesIO()
        torch.save(x_fp8, byte_stream)
        x_bytes = byte_stream.getvalue()

        # Mess up and delete old tensor
        x_fp8._data.zero_()
        x_fp8._scale_inv.zero_()
        del x_fp8, byte_stream

        # Deserialize tensor
        x_fp8 = torch.load(io.BytesIO(x_bytes))
        del x_bytes

        # Check results
        tols = dict(rtol=0, atol=0)
        t.assert_close(x_fp8, x_ref, **tols)

        # Make sure we are not trivially passing tests
        x_fp8._data.zero_()
        x_fp8._scale_inv.zero_()
        with pytest.raises(AssertionError):
            t.assert_close(x_fp8, x_ref, **tols)
//  [python]
def test_serialization(
        self,
        dims: DimsType = [2,3,5],
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 0.5,
        dtype: torch.dtype = torch.float32,
    ):

        # Initialize random data
        dims = _to_list(dims)
        x_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_ref = x_fp8.from_float8()

        # Serialize tensor
        byte_stream = io.BytesIO()
        torch.save(x_fp8, byte_stream)
        x_bytes = byte_stream.getvalue()

        # Mess up and delete old tensor
        x_fp8._data.zero_()
        x_fp8._scale_inv.zero_()
        del x_fp8, byte_stream

        # Deserialize tensor
        x_fp8 = torch.load(io.BytesIO(x_bytes))
        del x_bytes

        # Check results
        tols = dict(rtol=0, atol=0)
        torch.testing.assert_close(x_fp8, x_ref, **tols)

        # Make sure we are not trivially passing tests
        x_fp8._data.zero_()
        x_fp8._scale_inv.zero_()
        with pytest.raises(AssertionError):
            torch.testing.assert_close(x_fp8, x_ref, **tols)
//  [python]
def torch_assert_close(actual, expected, **kwargs):
    # assert_close was added around pt-1.9, it does better checks - e.g will check dimensions match
    if hasattr(torch.testing, "assert_close"):
        return torch.testing.assert_close(actual, expected, **kwargs)
    else:
        kwargs.pop("msg", None) # doesn't have msg arg
        return torch.allclose(actual, expected, **kwargs)
//  [python]
def test_basic_ops(
        self,
        dims: DimsType = 23,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 3.5,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        """Test basic out-of-place ops"""

        # Initialize random data
        dims = _to_list(dims)
        x_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        y_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        y_fp8 = Float8Tensor.to_float8(
            y_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_ref = x_fp8.from_float8()
        y_ref = y_fp8.from_float8()

        # Exact operations
        t.assert_close(-x_fp8, -x_ref, rtol=0, atol=0)
        t.assert_close(x_fp8.abs(), x_ref.abs(), rtol=0, atol=0)

        # Operations with numerical error
        tols = _tols[fp8_dtype]
        t.assert_close(x_fp8 + y_fp8, x_ref + y_ref, **tols)
        t.assert_close(x_fp8 - y_fp8, x_ref - y_ref, **tols)
        t.assert_close(x_fp8 * y_fp8, x_ref * y_ref, **tols)
        t.assert_close(x_fp8 + y_ref, x_ref + y_ref, **tols)
        t.assert_close(x_ref + y_fp8, x_ref + y_ref, **tols)
        t.assert_close(torch.sin(x_fp8), torch.sin(x_ref), **tols)

        # Make sure we are not trivially passing tests
        with pytest.raises(AssertionError):
            t.assert_close(x_fp8 + y_fp8, x_ref - y_fp8, **tols)
//  [python]
def test_basic_ops(
        self,
        dims: DimsType = 23,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 3.5,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        """Test basic out-of-place ops"""

        # Initialize random data
        dims = _to_list(dims)
        x_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        y_ref = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        y_fp8 = Float8Tensor.to_float8(
            y_ref,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x_ref = x_fp8.from_float8()
        y_ref = y_fp8.from_float8()

        # Exact operations
        torch.testing.assert_close(-x_fp8, -x_ref, rtol=0, atol=0)
        torch.testing.assert_close(x_fp8.abs(), x_ref.abs(), rtol=0, atol=0)

        # Operations with numerical error
        tols = _tols[fp8_dtype]
        torch.testing.assert_close(x_fp8 + y_fp8, x_ref + y_ref, **tols)
        torch.testing.assert_close(x_fp8 - y_fp8, x_ref - y_ref, **tols)
        torch.testing.assert_close(x_fp8 * y_fp8, x_ref * y_ref, **tols)
        torch.testing.assert_close(x_fp8 + y_ref, x_ref + y_ref, **tols)
        torch.testing.assert_close(x_ref + y_fp8, x_ref + y_ref, **tols)
        torch.testing.assert_close(torch.sin(x_fp8), torch.sin(x_ref), **tols)

        # Make sure we are not trivially passing tests
        with pytest.raises(AssertionError):
            torch.testing.assert_close(x_fp8 + y_fp8, x_ref - y_fp8, **tols)
//  [python]
def test_transpose(
        self,
        dims: DimsType,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 0.5,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        """Test transpose"""

        # Initialize random data
        dims = _to_list(dims)
        x = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x = x_fp8.from_float8()

        # Perform transpose
        x_fp8_t = x_fp8.transpose_2d()
        x_t = x.transpose(0, 1)
        x_fp8_t = Float8Tensor.make_like(x_fp8, data=x_fp8_t)

        # Check results
        tols = dict(rtol=0, atol=0)
        t.assert_close(x_fp8_t, x_t, **tols)

        # Make sure we are not trivially passing the test
        with pytest.raises(AssertionError):
            t.assert_close(x_fp8_t, x, **tols)

        # Caching test.
        assert x_fp8._transpose_invalid, "Transpose cache must be invalid when not caching."
        x_fp8 += 0.5
        x = x_fp8.from_float8()
        x_fp8_t = Float8Tensor.make_like(x_fp8, data=x_fp8.transpose_2d(cache=True))
        x_t = x.transpose(0, 1)
        t.assert_close(x_fp8_t, x_t, **tols)
        assert not x_fp8._transpose_invalid, "Transpose cache reset incorrectly."

        # Inplace update test.
        x_fp8 += 0.5
        assert x_fp8._transpose_invalid, "Transpose cache not invalidated properly."
        x = x_fp8.from_float8()
        x_fp8_t = Float8Tensor.make_like(x_fp8, data=x_fp8.transpose_2d(cache=True))
        x_t = x.transpose(0, 1)
        t.assert_close(x_fp8_t, x_t, **tols)
        assert not x_fp8._transpose_invalid, "Transpose cache reset incorrectly."
//  [python]
def test_transpose(
        self,
        dims: DimsType,
        fp8_dtype: tex.DType = tex.DType.kFloat8E4M3,
        scale: float = 0.5,
        dtype: torch.dtype = torch.float32,
    ) -> None:
        """Test transpose"""

        # Initialize random data
        dims = _to_list(dims)
        x = 2 * torch.rand(dims, dtype=dtype, device="cpu") - 1
        x_fp8 = Float8Tensor.to_float8(
            x,
            fp8_dtype=fp8_dtype,
            scale=torch.full([1], scale),
        )
        x = x_fp8.from_float8()

        # Perform transpose
        x_fp8_t = x_fp8.transpose_2d()
        x_t = x.transpose(0, 1)
        x_fp8_t = Float8Tensor.make_like(x_fp8, data=x_fp8_t)

        # Check results
        tols = dict(rtol=0, atol=0)
        torch.testing.assert_close(x_fp8_t, x_t, **tols)

        # Make sure we are not trivially passing the test
        with pytest.raises(AssertionError):
            torch.testing.assert_close(x_fp8_t, x, **tols)

        # Caching test.
        assert x_fp8._transpose_invalid, "Transpose cache must be invalid when not caching."
        x_fp8 += 0.5
        x = x_fp8.from_float8()
        x_fp8_t = Float8Tensor.make_like(x_fp8, data=x_fp8.transpose_2d(cache=True))
        x_t = x.transpose(0, 1)
        torch.testing.assert_close(x_fp8_t, x_t, **tols)
        assert not x_fp8._transpose_invalid, "Transpose cache reset incorrectly."

        # Inplace update test.
        x_fp8 += 0.5
        assert x_fp8._transpose_invalid, "Transpose cache not invalidated properly."
        x = x_fp8.from_float8()
        x_fp8_t = Float8Tensor.make_like(x_fp8, data=x_fp8.transpose_2d(cache=True))
        x_t = x.transpose(0, 1)
        torch.testing.assert_close(x_fp8_t, x_t, **tols)
        assert not x_fp8._transpose_invalid, "Transpose cache reset incorrectly."
//  [python]
def test_fill_unfill(batch_size):
    """Test that filling and unfilling results in the same tensor"""
    matrix = torch.randn(batch_size + (DIM, DIM))
    lower_vec = unfill_triangular(matrix, upper=False)
    upper_vec = unfill_triangular(matrix, upper=True)

    t.assert_close(torch.tril(matrix, diagonal=-1), fill_triangular(lower_vec, upper=False))
    t.assert_close(torch.triu(matrix, diagonal=1), fill_triangular(upper_vec, upper=True))
    t.assert_close(
        torch.tril(matrix, diagonal=-1).transpose(-2, -1), fill_triangular(lower_vec, upper=True)
    )
    t.assert_close(
        torch.triu(matrix, diagonal=1).transpose(-2, -1), fill_triangular(upper_vec, upper=False)
    )
//  [python]
def test_fill_unfill(batch_size):
    """Test that filling and unfilling results in the same tensor"""
    matrix = torch.randn(batch_size + (DIM, DIM))
    lower_vec = unfill_triangular(matrix, upper=False)
    upper_vec = unfill_triangular(matrix, upper=True)

    torch.testing.assert_close(torch.tril(matrix, diagonal=-1), fill_triangular(lower_vec, upper=False))
    torch.testing.assert_close(torch.triu(matrix, diagonal=1), fill_triangular(upper_vec, upper=True))
    torch.testing.assert_close(
        torch.tril(matrix, diagonal=-1).transpose(-2, -1), fill_triangular(lower_vec, upper=True)
    )
    torch.testing.assert_close(
        torch.triu(matrix, diagonal=1).transpose(-2, -1), fill_triangular(upper_vec, upper=False)
    )
//  [python]
def test_integration(self):
        # fmt: off
        EXPECTED_INPUT_FEATURES = torch.tensor(
            [
                -5.0229, -6.1358, -5.8346, -5.4447, -5.6707, -5.8577, -5.0464, -5.0058,
                -5.6015, -5.6410, -5.4325, -5.6116, -5.3700, -5.7956, -5.3196, -5.3274,
                -5.9655, -5.6057, -5.8382, -5.9602, -5.9005, -5.9123, -5.7669, -6.1441,
                -5.5168, -5.1405, -5.3927, -6.0032, -5.5784, -5.3728
            ],
        )
        # fmt: on

        input_speech, sr = self._load_datasamples(1)

        feature_extractor = UnivNetFeatureExtractor()
        input_features = feature_extractor(input_speech, sampling_rate=sr[0], return_tensors="pt").input_features
        self.assertEqual(input_features.shape, (1, 548, 100))

        input_features_mean = torch.mean(input_features)
        input_features_stddev = torch.std(input_features)

        EXPECTED_MEAN = torch.tensor(-6.18862009)
        EXPECTED_STDDEV = torch.tensor(2.80845642)

        self.assertTrue(torch.allclose(input_features_mean, EXPECTED_MEAN, atol=5e-5, rtol=5e-6))
        self.assertTrue(torch.allclose(input_features_stddev, EXPECTED_STDDEV))
        self.assertTrue(torch.allclose(input_features[0, :30, 0], EXPECTED_INPUT_FEATURES, atol=1e-4, rtol=1e-5))
//  [python]
def test_integration(self):
        # fmt: off
        EXPECTED_INPUT_FEATURES = torch.tensor(
            [
                -5.0229, -6.1358, -5.8346, -5.4447, -5.6707, -5.8577, -5.0464, -5.0058,
                -5.6015, -5.6410, -5.4325, -5.6116, -5.3700, -5.7956, -5.3196, -5.3274,
                -5.9655, -5.6057, -5.8382, -5.9602, -5.9005, -5.9123, -5.7669, -6.1441,
                -5.5168, -5.1405, -5.3927, -6.0032, -5.5784, -5.3728
            ],
        )
        # fmt: on

        input_speech, sr = self._load_datasamples(1)

        feature_extractor = UnivNetFeatureExtractor()
        input_features = feature_extractor(input_speech, sampling_rate=sr[0], return_tensors="pt").input_features
        self.assertEqual(input_features.shape, (1, 548, 100))

        input_features_mean = torch.mean(input_features)
        input_features_stddev = torch.std(input_features)

        EXPECTED_MEAN = torch.tensor(-6.18862009)
        EXPECTED_STDDEV = torch.tensor(2.80845642)

        torch.testing.assert_close(input_features_mean, EXPECTED_MEAN, atol=5e-5, rtol=5e-6)
        torch.testing.assert_close(input_features_stddev, EXPECTED_STDDEV)
        torch.testing.assert_close(input_features[0, :30, 0], EXPECTED_INPUT_FEATURES, atol=1e-4, rtol=1e-5)
//  [python]
def test_unfill_fill(batch_size):
    """Test that unfilling and filling results in the same tensor"""
    vec = torch.randn(batch_size + (DIM * (DIM - 1) // 2,))
    lower_mat = fill_triangular(vec, upper=False)
    upper_mat = fill_triangular(vec, upper=True)

    t.assert_close(unfill_triangular(lower_mat, upper=False), vec)
    t.assert_close(unfill_triangular(upper_mat, upper=True), vec)
    t.assert_close(unfill_triangular(lower_mat, upper=True), torch.zeros_like(vec))
    t.assert_close(unfill_triangular(upper_mat, upper=False), torch.zeros_like(vec))
//  [python]
def test_unfill_fill(batch_size):
    """Test that unfilling and filling results in the same tensor"""
    vec = torch.randn(batch_size + (DIM * (DIM - 1) // 2,))
    lower_mat = fill_triangular(vec, upper=False)
    upper_mat = fill_triangular(vec, upper=True)

    torch.testing.assert_close(unfill_triangular(lower_mat, upper=False), vec)
    torch.testing.assert_close(unfill_triangular(upper_mat, upper=True), vec)
    torch.testing.assert_close(unfill_triangular(lower_mat, upper=True), torch.zeros_like(vec))
    torch.testing.assert_close(unfill_triangular(upper_mat, upper=False), torch.zeros_like(vec))
//  [python]
def _test_binary_binned_auroc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int,
        threshold: Union[torch.Tensor, int],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        my_compute_result = binary_binned_auroc(
            input,
            target,
            num_tasks=num_tasks,
            threshold=threshold,
        )
        t.assert_allclose(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_binary_binned_auroc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_tasks: int,
        threshold: Union[torch.Tensor, int],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
    ) -> None:
        my_compute_result = binary_binned_auroc(
            input,
            target,
            num_tasks=num_tasks,
            threshold=threshold,
        )
        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def test_pad_sequence_from_left():
    sequences = [
        torch.tensor([0.0, 1.0, 2.0]),
        torch.tensor(
            [
                3.0,
            ]
        ),
        torch.tensor(
            [
                6.0,
                7.0,
            ]
        ),
    ]
    expected = torch.tensor([[0.0, 1.0, 2.0], [-1.0, -1.0, 3.0], [-1.0, 6.0, 7.0]])
    actual = torch_ops.pad_sequence_from_left(sequences, batch_first=True, padding_value=-1)
    torch.testing.assert_allclose(actual, expected)
//  [python]
def test_pad_sequence_from_left():
    sequences = [
        torch.tensor([0.0, 1.0, 2.0]),
        torch.tensor(
            [
                3.0,
            ]
        ),
        torch.tensor(
            [
                6.0,
                7.0,
            ]
        ),
    ]
    expected = torch.tensor([[0.0, 1.0, 2.0], [-1.0, -1.0, 3.0], [-1.0, 6.0, 7.0]])
    actual = torch_ops.pad_sequence_from_left(sequences, batch_first=True, padding_value=-1)
    torch.testing.assert_close(actual, expected)
//  [python]
def _test_multiclass_binned_auroc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        threshold: Union[torch.Tensor, int],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
        average: Optional[str] = "macro",
    ) -> None:
        my_compute_result = multiclass_binned_auroc(
            input,
            target,
            num_classes=num_classes,
            threshold=threshold,
            average=average,
        )
        t.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def _test_multiclass_binned_auroc_with_input(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        num_classes: int,
        threshold: Union[torch.Tensor, int],
        compute_result: Tuple[torch.Tensor, torch.Tensor],
        average: Optional[str] = "macro",
    ) -> None:
        my_compute_result = multiclass_binned_auroc(
            input,
            target,
            num_classes=num_classes,
            threshold=threshold,
            average=average,
        )
        torch.testing.assert_close(
            my_compute_result,
            compute_result,
            equal_nan=True,
            atol=1e-8,
            rtol=1e-5,
        )
//  [python]
def testAdd(self):
        device = get_test_device()
        x = torch.randn(2, 3, 4, device=device)
        y = torch.randn(2, 3, 4, device=device)
        z = torch.zeros(2, 3, 4, device=device)

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        y_lazy = y.detach().clone().to(device=device)
        z_lazy = z.detach().clone().to(device=device)

        for i in range(10):
            z += x + y

        for i in range(10):
            z_lazy += x_lazy + y_lazy
            torch._lazy.mark_step()

        torch.testing.assert_allclose(z.cpu(), z_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 14
        metrics.reset()
        torch._lazy.ir_cache.reset()
//  [python]
def testAdd(self):
        device = get_test_device()
        x = torch.randn(2, 3, 4, device=device)
        y = torch.randn(2, 3, 4, device=device)
        z = torch.zeros(2, 3, 4, device=device)

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        y_lazy = y.detach().clone().to(device=device)
        z_lazy = z.detach().clone().to(device=device)

        for i in range(10):
            z += x + y

        for i in range(10):
            z_lazy += x_lazy + y_lazy
            torch._lazy.mark_step()

        torch.testing.assert_close(z.cpu(), z_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 14
        metrics.reset()
        torch._lazy.ir_cache.reset()
//  [python]
print("SO3 identity", i1, i1.shape)
print("SE3 identity", i2, i2.shape)

# Indexing
g1_slice = g1[2:4]
assert g1_slice.shape == (2, 3, 4)
torch.allclose(g1_slice._t, g1._t[2:4])
//  [python]
print("SO3 identity", i1, i1.shape)
print("SE3 identity", i2, i2.shape)

# Indexing
g1_slice = g1[2:4]
assert g1_slice.shape == (2, 3, 4)
torch.testing.assert_close(g1_slice._t, g1._t[2:4])
//  [python]
def testAddSub(self):
        device = get_test_device()
        x = torch.randn(2, 3, 4, device=device)
        y = torch.randn(2, 3, 4, device=device)
        z = torch.zeros(2, 3, 4, device=device)

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        y_lazy = y.detach().clone().to(device=device)
        z_lazy = z.detach().clone().to(device=device)

        for i in range(10):
            if i < 5:
                z += x + y
            else:
                z += x - y

        for i in range(10):
            if i < 5:
                z_lazy += x_lazy + y_lazy
            else:
                z_lazy += x_lazy - y_lazy
            torch._lazy.mark_step()

        torch.testing.assert_allclose(z.cpu(), z_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 8
        metrics.reset()
        torch._lazy.ir_cache.reset()
//  [python]
def testAddSub(self):
        device = get_test_device()
        x = torch.randn(2, 3, 4, device=device)
        y = torch.randn(2, 3, 4, device=device)
        z = torch.zeros(2, 3, 4, device=device)

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        y_lazy = y.detach().clone().to(device=device)
        z_lazy = z.detach().clone().to(device=device)

        for i in range(10):
            if i < 5:
                z += x + y
            else:
                z += x - y

        for i in range(10):
            if i < 5:
                z_lazy += x_lazy + y_lazy
            else:
                z_lazy += x_lazy - y_lazy
            torch._lazy.mark_step()

        torch.testing.assert_close(z.cpu(), z_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 8
        metrics.reset()
        torch._lazy.ir_cache.reset()
//  [python]
def testAddSubFallback(self):
        torch._lazy.config.set_force_fallback("aten::sub")
        device = get_test_device()
        x = torch.randn(2, 3, 4, device=device)
        y = torch.randn(2, 3, 4, device=device)
        z = torch.zeros(2, 3, 4, device=device)

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        y_lazy = y.detach().clone().to(device=device)
        z_lazy = z.detach().clone().to(device=device)

        for i in range(10):
            if i < 5:
                z += x + y
            else:
                z += x - y

        for i in range(10):
            if i < 5:
                z_lazy += x_lazy + y_lazy
            else:
                z_lazy += x_lazy - y_lazy
            torch._lazy.mark_step()

        self.assertEqual(z.cpu(), z_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 8
        metrics.reset()
        torch._lazy.ir_cache.reset()
        torch._lazy.config.set_force_fallback("")
//  [python]
def testAddSubFallback(self):
        torch._lazy.config.set_force_fallback("aten::sub")
        device = get_test_device()
        x = torch.randn(2, 3, 4, device=device)
        y = torch.randn(2, 3, 4, device=device)
        z = torch.zeros(2, 3, 4, device=device)

        device = "lazy"
        x_lazy = x.detach().clone().to(device=device)
        y_lazy = y.detach().clone().to(device=device)
        z_lazy = z.detach().clone().to(device=device)

        for i in range(10):
            if i < 5:
                z += x + y
            else:
                z += x - y

        for i in range(10):
            if i < 5:
                z_lazy += x_lazy + y_lazy
            else:
                z_lazy += x_lazy - y_lazy
            torch._lazy.mark_step()

        torch.testing.assert_close(z.cpu(), z_lazy.cpu())
        assert metrics.counter_value("IrNodeReused_torch::lazy::AddTensor") >= 8
        metrics.reset()
        torch._lazy.ir_cache.reset()
        torch._lazy.config.set_force_fallback("")
//  [python]
def test_sum_reduce_to_one(
        self,
        # pyre-fixme[2]: Parameter must be annotated.
        num_inputs,
        # pyre-fixme[2]: Parameter must be annotated.
        num_gpus,
        # pyre-fixme[2]: Parameter must be annotated.
        r,
    ) -> None:
        dst_device = torch.device(f"cuda:{r.randint(0, num_gpus - 1)}")
        with torch.cuda.device(dst_device):
            inputs = [torch.randn(10, 20) for _ in range(num_inputs)]
            cuda_inputs = [
                input.to(f"cuda:{i % num_gpus}") for i, input in enumerate(inputs)
            ]
            cuda_output = torch.ops.fbgemm.sum_reduce_to_one(cuda_inputs, dst_device)
            self.assertEqual(cuda_output.device, dst_device)
            torch.testing.assert_allclose(
                cuda_output.cpu(), torch.stack(inputs).sum(dim=0)
            )
//  [python]
def test_sum_reduce_to_one(
        self,
        # pyre-fixme[2]: Parameter must be annotated.
        num_inputs,
        # pyre-fixme[2]: Parameter must be annotated.
        num_gpus,
        # pyre-fixme[2]: Parameter must be annotated.
        r,
    ) -> None:
        dst_device = torch.device(f"cuda:{r.randint(0, num_gpus - 1)}")
        with torch.cuda.device(dst_device):
            inputs = [torch.randn(10, 20) for _ in range(num_inputs)]
            cuda_inputs = [
                input.to(f"cuda:{i % num_gpus}") for i, input in enumerate(inputs)
            ]
            cuda_output = torch.ops.fbgemm.sum_reduce_to_one(cuda_inputs, dst_device)
            self.assertEqual(cuda_output.device, dst_device)
            torch.testing.assert_close(
                cuda_output.cpu(), torch.stack(inputs).sum(dim=0)
            )
//  [python]
def test_all_to_one_device(
        self,
        # pyre-fixme[2]: Parameter must be annotated.
        num_inputs,
        # pyre-fixme[2]: Parameter must be annotated.
        num_gpus,
        # pyre-fixme[2]: Parameter must be annotated.
        r,
    ) -> None:
        dst_device = torch.device(f"cuda:{r.randint(0, num_gpus - 1)}")
        with torch.cuda.device(dst_device):
            inputs = [torch.randn(10, 20) for _ in range(num_inputs)]
            cuda_inputs = [
                input.to(f"cuda:{i % num_gpus}") for i, input in enumerate(inputs)
            ]
            cuda_outputs = torch.ops.fbgemm.all_to_one_device(cuda_inputs, dst_device)
            for i, o in zip(inputs, cuda_outputs):
                self.assertEqual(o.device, dst_device)
                self.assertEqual(o.cpu(), i)
//  [python]
def test_all_to_one_device(
        self,
        # pyre-fixme[2]: Parameter must be annotated.
        num_inputs,
        # pyre-fixme[2]: Parameter must be annotated.
        num_gpus,
        # pyre-fixme[2]: Parameter must be annotated.
        r,
    ) -> None:
        dst_device = torch.device(f"cuda:{r.randint(0, num_gpus - 1)}")
        with torch.cuda.device(dst_device):
            inputs = [torch.randn(10, 20) for _ in range(num_inputs)]
            cuda_inputs = [
                input.to(f"cuda:{i % num_gpus}") for i, input in enumerate(inputs)
            ]
            cuda_outputs = torch.ops.fbgemm.all_to_one_device(cuda_inputs, dst_device)
            for i, o in zip(inputs, cuda_outputs):
                self.assertEqual(o.device, dst_device)
                torch.testing.assert_close(o.cpu(), i)