jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_value_iteration():
    environment = EasyGridWorld()
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = value_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_close(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )
    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(environment, GAMMA, value_function, OPTIMAL_POLICY, pred_p)

    environment = EasyGridWorld(terminal_states=[22])
    GAMMA = 0.9
    EPS = 1e-3
    policy, value_function = value_iteration(environment, GAMMA, eps=EPS)

    torch.testing.assert_close(
        value_function.table,
        torch.tensor([OPTIMAL_VALUE_WITH_TERMINAL]).unsqueeze(-1),
        atol=0.05,
        rtol=EPS,
    )

    pred_p = policy.table.argmax(dim=0)
    assert_policy_equality(
        environment, GAMMA, value_function, OPTIMAL_POLICY_WITH_TERMINAL, pred_p
    )
//  [python]
def test_masked_load(dtype_str, size, size_diff, device='cuda'):
    dtype = getattr(torch, dtype_str)
    check_type_supported(dtype)  # bfloat16 on cc < 80 will not be tested

    input_size = size - size_diff
    output_size = size
    if dtype_str == 'bool':
        input = torch.randint(0, 2, (input_size,), dtype=dtype, device=device)
    elif dtype_str in int_dtypes or dtype_str in uint_dtypes:
        input = torch.randint(0, 127, (input_size,), dtype=dtype, device=device)
    else:
        input = torch.rand(input_size, dtype=dtype, device=device)
    output = torch.zeros((output_size,), dtype=dtype, device=device)

    @triton.jit
    def _kernel(in_ptr, out_ptr, in_size: tl.constexpr, out_size: tl.constexpr):
        in_offsets = tl.arange(0, out_size)
        # Load inputs.
        x = GENERATE_TEST_HERE
        # Store output
        output_offsets = tl.arange(0, out_size)
        tl.store(out_ptr + output_offsets, x)

    mask_str = "mask=in_offsets < in_size, other=1" if size_diff > 0 else "None"
    kernel = patch_kernel(_kernel, {'GENERATE_TEST_HERE': f"tl.load(in_ptr + in_offsets, {mask_str})"})
    kernel[(1,)](input, output, input_size, output_size)

    reference_out = torch.cat((input, torch.ones((size_diff,), dtype=dtype, device=device)))
    torch.testing.assert_allclose(output, reference_out)
//  [python]
def test_masked_load(dtype_str, size, size_diff, device='cuda'):
    dtype = getattr(torch, dtype_str)
    check_type_supported(dtype)  # bfloat16 on cc < 80 will not be tested

    input_size = size - size_diff
    output_size = size
    if dtype_str == 'bool':
        input = torch.randint(0, 2, (input_size,), dtype=dtype, device=device)
    elif dtype_str in int_dtypes or dtype_str in uint_dtypes:
        input = torch.randint(0, 127, (input_size,), dtype=dtype, device=device)
    else:
        input = torch.rand(input_size, dtype=dtype, device=device)
    output = torch.zeros((output_size,), dtype=dtype, device=device)

    @triton.jit
    def _kernel(in_ptr, out_ptr, in_size: tl.constexpr, out_size: tl.constexpr):
        in_offsets = tl.arange(0, out_size)
        # Load inputs.
        x = GENERATE_TEST_HERE
        # Store output
        output_offsets = tl.arange(0, out_size)
        tl.store(out_ptr + output_offsets, x)

    mask_str = "mask=in_offsets < in_size, other=1" if size_diff > 0 else "None"
    kernel = patch_kernel(_kernel, {'GENERATE_TEST_HERE': f"tl.load(in_ptr + in_offsets, {mask_str})"})
    kernel[(1,)](input, output, input_size, output_size)

    reference_out = torch.cat((input, torch.ones((size_diff,), dtype=dtype, device=device)))
    torch.testing.assert_close(output, reference_out)
//  [python]
def test_region_grow(self):
        cluster_idx = region_grow(self.pos, self.labels, self.batch, radius=2, min_cluster_size=1)
        self.assertEqual(len(cluster_idx), 6)
        torch.testing.assert_allclose(cluster_idx[0], torch.tensor([0, 1]))
        torch.testing.assert_allclose(cluster_idx[1], torch.tensor([4]))
        torch.testing.assert_allclose(cluster_idx[2], torch.tensor([2]))
        torch.testing.assert_allclose(cluster_idx[3], torch.tensor([3]))
        torch.testing.assert_allclose(cluster_idx[4], torch.tensor([5, 6]))
        torch.testing.assert_allclose(cluster_idx[5], torch.tensor([7]))
//  [python]
def test_region_grow(self):
        cluster_idx = region_grow(self.pos, self.labels, self.batch, radius=2, min_cluster_size=1)
        self.assertEqual(len(cluster_idx), 6)
        torch.testing.assert_close(cluster_idx[0], torch.tensor([0, 1]))
        torch.testing.assert_close(cluster_idx[1], torch.tensor([4]))
        torch.testing.assert_close(cluster_idx[2], torch.tensor([2]))
        torch.testing.assert_close(cluster_idx[3], torch.tensor([3]))
        torch.testing.assert_close(cluster_idx[4], torch.tensor([5, 6]))
        torch.testing.assert_close(cluster_idx[5], torch.tensor([7]))
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_allclose(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_close(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)
//  [python]
def test_load_save(model_name: str, device: torch.device):
    model, config = get_model_config(
        model_name, learning_rate=1e-4, warmup_steps=1000, weight_decay=0.0
    )
    optimizer, lr_scheduler = load_or_construct_optimizer_lr_scheduler(model, config)
    with TemporaryDirectory() as dirname:
        save_model(
            os.path.join(dirname, "best_model.pt"),
            model,
            optimizer,
            lr_scheduler,
            training_config=config,
        )

        model2_wrapped = load_model(dirname)
        assert model2_wrapped.loss is None
        model2 = unwrap_loss_wrapper(model2_wrapped)
        model2_params = dict(model2.named_parameters())
        for name, param in model.named_parameters():
            assert name in model2_params
            torch.testing.assert_allclose(model2_params[name], param, rtol=0, atol=0)

        optimizer2, lr_scheduler2 = load_or_construct_optimizer_lr_scheduler(
            model2, get_training_config(saved_model_path=dirname)
        )
        assert optimizer.state_dict() == optimizer2.state_dict()
        assert lr_scheduler.state_dict() == lr_scheduler2.state_dict()

    # Wrapped in `LossWrapper`
    loss = NegativeLogLikelihoodLoss()
    model_wrapped = LossWrapper(module=model, loss=loss)
    with TemporaryDirectory() as dirname:
        save_model(
            os.path.join(dirname, "best_model.pt"),
            model_wrapped,
            optimizer,
            lr_scheduler,
            training_config=config,
        )

        model2_wrapped = load_model(dirname)
        assert model2_wrapped.loss is None
        # Check model parameters ignoring the loss.
        check_module_params(unwrap_loss_wrapper(model_wrapped), unwrap_loss_wrapper(model2_wrapped))
//  [python]
def test_load_save(model_name: str, device: torch.device):
    model, config = get_model_config(
        model_name, learning_rate=1e-4, warmup_steps=1000, weight_decay=0.0
    )
    optimizer, lr_scheduler = load_or_construct_optimizer_lr_scheduler(model, config)
    with TemporaryDirectory() as dirname:
        save_model(
            os.path.join(dirname, "best_model.pt"),
            model,
            optimizer,
            lr_scheduler,
            training_config=config,
        )

        model2_wrapped = load_model(dirname)
        assert model2_wrapped.loss is None
        model2 = unwrap_loss_wrapper(model2_wrapped)
        model2_params = dict(model2.named_parameters())
        for name, param in model.named_parameters():
            assert name in model2_params
            torch.testing.assert_close(model2_params[name], param, rtol=0, atol=0)

        optimizer2, lr_scheduler2 = load_or_construct_optimizer_lr_scheduler(
            model2, get_training_config(saved_model_path=dirname)
        )
        assert optimizer.state_dict() == optimizer2.state_dict()
        assert lr_scheduler.state_dict() == lr_scheduler2.state_dict()

    # Wrapped in `LossWrapper`
    loss = NegativeLogLikelihoodLoss()
    model_wrapped = LossWrapper(module=model, loss=loss)
    with TemporaryDirectory() as dirname:
        save_model(
            os.path.join(dirname, "best_model.pt"),
            model_wrapped,
            optimizer,
            lr_scheduler,
            training_config=config,
        )

        model2_wrapped = load_model(dirname)
        assert model2_wrapped.loss is None
        # Check model parameters ignoring the loss.
        check_module_params(unwrap_loss_wrapper(model_wrapped), unwrap_loss_wrapper(model2_wrapped))
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (1000, ))
        torch.testing.assert_allclose(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)
//  [python]
def test_predict(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]

        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)
        predictions = model.predict(inputs)
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))

        predictions = model.predict(inputs, data_samples)
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (1000, ))
        torch.testing.assert_close(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)
//  [python]
def tracer_tutorial():
	# Automatic trace checking.

	if False:
		def loop_in_traced_fn(x):
			result = x[0]
			for i in range(x.size(0)):
				result = result * x[i]
			return result

		inputs = (torch.rand(3, 4, 5),)
		check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

		# One way to automatically catch many errors in traces is by using check_inputs on the torch.jit.trace() API. 
		traced = torch.jit.trace(loop_in_traced_fn, inputs, check_inputs=check_inputs)

	if False:
		def fn(x):
			result = x[0]
			for i in range(x.size(0)):
				result = result * x[i]
			return result

		inputs = (torch.rand(3, 4, 5),)
		check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

		scripted_fn = torch.jit.script(fn)
		print(scripted_fn.graph)
		#print(str(scripted_fn.graph).strip())

		# Data-dependent control flow can be captured using torch.jit.script().
		for input_tuple in [inputs] + check_inputs:
			torch.testing.assert_allclose(fn(*input_tuple), scripted_fn(*input_tuple))

	#--------------------
	# Tracer warnings.

	if True:
		def fill_row_zero(x):
			#x[0] = torch.rand(*x.shape[1:2])  # An in-place assignment on a slice (a view) of a Tensor.
			x = torch.cat((torch.rand(1, *x.shape[1:2]), x[1:2]), dim=0)  # The code to not use the in-place update. The result tensor out-of-place with torch.cat.
			return x

		traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
		print(traced.graph)
//  [python]
def tracer_tutorial():
	# Automatic trace checking.

	if False:
		def loop_in_traced_fn(x):
			result = x[0]
			for i in range(x.size(0)):
				result = result * x[i]
			return result

		inputs = (torch.rand(3, 4, 5),)
		check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

		# One way to automatically catch many errors in traces is by using check_inputs on the torch.jit.trace() API. 
		traced = torch.jit.trace(loop_in_traced_fn, inputs, check_inputs=check_inputs)

	if False:
		def fn(x):
			result = x[0]
			for i in range(x.size(0)):
				result = result * x[i]
			return result

		inputs = (torch.rand(3, 4, 5),)
		check_inputs = [(torch.rand(4, 5, 6),), (torch.rand(2, 3, 4),)]

		scripted_fn = torch.jit.script(fn)
		print(scripted_fn.graph)
		#print(str(scripted_fn.graph).strip())

		# Data-dependent control flow can be captured using torch.jit.script().
		for input_tuple in [inputs] + check_inputs:
			torch.testing.assert_close(fn(*input_tuple), scripted_fn(*input_tuple))

	#--------------------
	# Tracer warnings.

	if True:
		def fill_row_zero(x):
			#x[0] = torch.rand(*x.shape[1:2])  # An in-place assignment on a slice (a view) of a Tensor.
			x = torch.cat((torch.rand(1, *x.shape[1:2]), x[1:2]), dim=0)  # The code to not use the in-place update. The result tensor out-of-place with torch.cat.
			return x

		traced = torch.jit.trace(fill_row_zero, (torch.rand(3, 4),))
		print(traced.graph)
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]
        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

        # test forward train
        losses = model(inputs, data_samples, mode='loss')
        self.assertGreater(losses['loss'].item(), 0)

        # test forward test
        predictions = model(inputs, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model(inputs, data_samples, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_allclose(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]
        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

        # test forward train
        losses = model(inputs, data_samples, mode='loss')
        self.assertGreater(losses['loss'].item(), 0)

        # test forward test
        predictions = model(inputs, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))

        predictions = model(inputs, data_samples, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (10, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (10, ))
        torch.testing.assert_close(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def forward(device):
    make_cudnn_deterministic()
    model_dim = 8
    num_experts = dist.get_world_size(dist.group.WORLD)
    input = torch.randn(4, 16, model_dim).to(device)
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use identity matrix
    expert.weight = torch.nn.Parameter(torch.eye(model_dim))
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # Re-assembled output should match input due to identity expert.
    torch.testing.assert_allclose(input, output)
//  [python]
def forward(device):
    make_cudnn_deterministic()
    model_dim = 8
    num_experts = dist.get_world_size(dist.group.WORLD)
    input = torch.randn(4, 16, model_dim).to(device)
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use identity matrix
    expert.weight = torch.nn.Parameter(torch.eye(model_dim))
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # Re-assembled output should match input due to identity expert.
    torch.testing.assert_close(input, output)
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]
        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

        # test forward train
        losses = model(inputs, data_samples, mode='loss')
        self.assertGreater(losses['loss'].item(), 0)

        # test forward test
        predictions = model(inputs, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))

        predictions = model(inputs, data_samples, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (1000, ))
        torch.testing.assert_allclose(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def test_forward(self):
        inputs = torch.rand(1, 3, 224, 224)
        data_samples = [ClsDataSample().set_gt_label(1)]
        model: ImageClassifier = MODELS.build(self.DEFAULT_ARGS)

        # test pure forward
        outs = model(inputs)
        self.assertIsInstance(outs, torch.Tensor)

        # test forward train
        losses = model(inputs, data_samples, mode='loss')
        self.assertGreater(losses['loss'].item(), 0)

        # test forward test
        predictions = model(inputs, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))

        predictions = model(inputs, data_samples, mode='predict')
        self.assertEqual(predictions[0].pred_label.score.shape, (1000, ))
        self.assertEqual(data_samples[0].pred_label.score.shape, (1000, ))
        torch.testing.assert_close(data_samples[0].pred_label.score,
                                      predictions[0].pred_label.score)

        # test forward with invalid mode
        with self.assertRaisesRegex(RuntimeError, 'Invalid mode "unknown"'):
            model(inputs, mode='unknown')
//  [python]
def forward_routing(device):
    make_cudnn_deterministic()
    model_dim = 8
    num_experts = dist.get_world_size()
    input = torch.randn(4, 16, model_dim).to(device)
    gate = RoundRobinGate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use scaling matrix (each rank has a different scale)
    scale = dist.get_rank() + 1
    expert.weight = torch.nn.Parameter(torch.eye(model_dim) * scale)
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # Verify that each token was sent to the correct expert by checking its scale.
    t = input.shape[1]
    for i in range(t):
        expert = i % num_experts
        torch.testing.assert_allclose(input[:, i] * (expert + 1), output[:, i])
//  [python]
def forward_routing(device):
    make_cudnn_deterministic()
    model_dim = 8
    num_experts = dist.get_world_size()
    input = torch.randn(4, 16, model_dim).to(device)
    gate = RoundRobinGate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use scaling matrix (each rank has a different scale)
    scale = dist.get_rank() + 1
    expert.weight = torch.nn.Parameter(torch.eye(model_dim) * scale)
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # Verify that each token was sent to the correct expert by checking its scale.
    t = input.shape[1]
    for i in range(t):
        expert = i % num_experts
        torch.testing.assert_close(input[:, i] * (expert + 1), output[:, i])
//  [python]
def backward(device):
    make_cudnn_deterministic()
    loss = torch.nn.MSELoss()
    model_dim = 8
    num_experts = dist.get_world_size(dist.group.WORLD)
    input = torch.randn(4, 16, model_dim).to(device)
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use identity matrix
    expert.weight = torch.nn.Parameter(torch.eye(model_dim))
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    output = loss(output, input)
    output.backward()
    torch.testing.assert_allclose(expert.weight.grad, torch.zeros_like(expert.weight))
//  [python]
def backward(device):
    make_cudnn_deterministic()
    loss = torch.nn.MSELoss()
    model_dim = 8
    num_experts = dist.get_world_size(dist.group.WORLD)
    input = torch.randn(4, 16, model_dim).to(device)
    gate = Top2Gate(model_dim, num_experts)
    expert = torch.nn.Linear(model_dim, model_dim, bias=False)
    # Use identity matrix
    expert.weight = torch.nn.Parameter(torch.eye(model_dim))
    moe = MOELayer(gate, expert).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    output = loss(output, input)
    output.backward()
    torch.testing.assert_close(expert.weight.grad, torch.zeros_like(expert.weight))
//  [python]
def forward_multi(device):
    make_cudnn_deterministic()
    torch.set_printoptions(threshold=5000)
    num_local_experts = 4
    model_dim = 4
    num_experts = dist.get_world_size(dist.group.WORLD) * num_local_experts
    input = torch.randn(4 * num_local_experts, 16, model_dim).to(device)
    gate = Top2Gate(model_dim, num_experts)
    experts = []
    for i in range(num_local_experts):
        expert = torch.nn.Linear(model_dim, model_dim, bias=False)
        # Use identity matrix
        expert.weight = torch.nn.Parameter(torch.eye(model_dim))
        experts += [expert]
    moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # 90% of the input should have gone to an expert
    assert (
        len(output.nonzero(as_tuple=False)) / output.numel() > 0.90
    ), f"{len(output.nonzero(as_tuple=False))} / {output.numel()}"
    # Except for zeros, re-assembled output should match input due to identity expert.
    torch.testing.assert_allclose(input, torch.where(output > 0, output, input))
//  [python]
def forward_multi(device):
    make_cudnn_deterministic()
    torch.set_printoptions(threshold=5000)
    num_local_experts = 4
    model_dim = 4
    num_experts = dist.get_world_size(dist.group.WORLD) * num_local_experts
    input = torch.randn(4 * num_local_experts, 16, model_dim).to(device)
    gate = Top2Gate(model_dim, num_experts)
    experts = []
    for i in range(num_local_experts):
        expert = torch.nn.Linear(model_dim, model_dim, bias=False)
        # Use identity matrix
        expert.weight = torch.nn.Parameter(torch.eye(model_dim))
        experts += [expert]
    moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # 90% of the input should have gone to an expert
    assert (
        len(output.nonzero(as_tuple=False)) / output.numel() > 0.90
    ), f"{len(output.nonzero(as_tuple=False))} / {output.numel()}"
    # Except for zeros, re-assembled output should match input due to identity expert.
    torch.testing.assert_close(input, torch.where(output > 0, output, input))
//  [python]
def forward_routing_multi(device):
    make_cudnn_deterministic()
    model_dim = 8
    num_local_experts = 4
    num_experts = dist.get_world_size(dist.group.WORLD) * num_local_experts
    input = torch.randn(4 * num_local_experts, 16, model_dim).to(device)
    gate = RoundRobinGate(model_dim, num_experts)
    experts = []
    for i in range(num_local_experts):
        expert = torch.nn.Linear(model_dim, model_dim, bias=False)
        # Use scaling matrix (each rank has a different scale)
        scale = dist.get_rank() * num_local_experts + i + 1
        expert.weight = torch.nn.Parameter(torch.eye(model_dim) * scale)
        experts += [expert]
    moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # Verify that each token was sent to the correct expert by checking its scale.
    t = input.shape[1]
    for i in range(t):
        expert = i % num_experts
        torch.testing.assert_allclose(input[:, i] * (expert + 1), output[:, i])
//  [python]
def forward_routing_multi(device):
    make_cudnn_deterministic()
    model_dim = 8
    num_local_experts = 4
    num_experts = dist.get_world_size(dist.group.WORLD) * num_local_experts
    input = torch.randn(4 * num_local_experts, 16, model_dim).to(device)
    gate = RoundRobinGate(model_dim, num_experts)
    experts = []
    for i in range(num_local_experts):
        expert = torch.nn.Linear(model_dim, model_dim, bias=False)
        # Use scaling matrix (each rank has a different scale)
        scale = dist.get_rank() * num_local_experts + i + 1
        expert.weight = torch.nn.Parameter(torch.eye(model_dim) * scale)
        experts += [expert]
    moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
    output = moe(input)
    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
    # Verify that each token was sent to the correct expert by checking its scale.
    t = input.shape[1]
    for i in range(t):
        expert = i % num_experts
        torch.testing.assert_close(input[:, i] * (expert + 1), output[:, i])
//  [python]
def test_resnetish(self):
        model = resnetish()
        model.eval()
        inputs = [self._prepare_image()]
        jit_out, pml_out = self._run_both(model, inputs)
        torch.testing.assert_allclose(jit_out, pml_out, rtol=0.01, atol=0.01)
        self._benchmark(model, inputs)
//  [python]
def test_resnetish(self):
        model = resnetish()
        model.eval()
        inputs = [self._prepare_image()]
        jit_out, pml_out = self._run_both(model, inputs)
        torch.testing.assert_close(jit_out, pml_out, rtol=0.01, atol=0.01)
        self._benchmark(model, inputs)
//  [python]
def test_autocast_fused_scale_mask_softmax(self):
        for dtype in autocast_dtypes:
            with self.subTest(f"{dtype}"):
                input_in_fp16 = dtype == torch.half
                input_in_bf16 = dtype == torch.bfloat16
                fused_fn, torch_fn = self._setup_fused_softmax(
                    input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding)

                attention_scores_0 = torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True)
                with torch.no_grad():
                    attention_scores_1 = attention_scores_0.clone().to(dtype).requires_grad_(True)
                mask = torch.randint(0, 2, (4, 1, 24, 24)).bool().cuda()

                expected = torch_fn(attention_scores_1, mask)
                with torch.cuda.amp.autocast(dtype=dtype):
                    actual = fused_fn(attention_scores_0, mask)
                    self.assertEqual(actual.dtype, dtype)
                torch.testing.assert_allclose(actual, expected)

                g0 = torch.rand_like(actual)
                with torch.no_grad():
                    g1 = g0.clone()
                expected.backward(g0)
                actual.backward(g1)
//  [python]
def test_autocast_fused_scale_mask_softmax(self):
        for dtype in autocast_dtypes:
            with self.subTest(f"{dtype}"):
                input_in_fp16 = dtype == torch.half
                input_in_bf16 = dtype == torch.bfloat16
                fused_fn, torch_fn = self._setup_fused_softmax(
                    input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding)

                attention_scores_0 = torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True)
                with torch.no_grad():
                    attention_scores_1 = attention_scores_0.clone().to(dtype).requires_grad_(True)
                mask = torch.randint(0, 2, (4, 1, 24, 24)).bool().cuda()

                expected = torch_fn(attention_scores_1, mask)
                with torch.cuda.amp.autocast(dtype=dtype):
                    actual = fused_fn(attention_scores_0, mask)
                    self.assertEqual(actual.dtype, dtype)
                torch.testing.assert_close(actual, expected)

                g0 = torch.rand_like(actual)
                with torch.no_grad():
                    g1 = g0.clone()
                expected.backward(g0)
                actual.backward(g1)
//  [python]
def test_autocast_fused_upper_triangle_mask_softmax(self):
        for dtype in autocast_dtypes:
            with self.subTest(f"{dtype}"):
                input_in_fp16 = dtype == torch.half
                input_in_bf16 = dtype == torch.bfloat16
                fused_fn, torch_fn = self._setup_fused_softmax(
                    input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.causal)

                attn_weights_0 = torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True)
                with torch.no_grad():
                    attn_weights_1 = attn_weights_0.clone().to(dtype).requires_grad_(True)
                total_mask = (~(
                    torch.tril(torch.randn((24, 24), device="cuda")).bool()
                ).unsqueeze(0).unsqueeze(0))

                with torch.cuda.amp.autocast(dtype=dtype):
                    actual = fused_fn(attn_weights_0, total_mask)
                    self.assertEqual(actual.dtype, dtype)
                expected = torch_fn(attn_weights_1, total_mask)
                torch.testing.assert_allclose(actual, expected)

                g0 = torch.randn_like(actual)
                with torch.no_grad():
                    g1 = g0.clone()
                actual.backward(g0)
                expected.backward(g1)
//  [python]
def test_autocast_fused_upper_triangle_mask_softmax(self):
        for dtype in autocast_dtypes:
            with self.subTest(f"{dtype}"):
                input_in_fp16 = dtype == torch.half
                input_in_bf16 = dtype == torch.bfloat16
                fused_fn, torch_fn = self._setup_fused_softmax(
                    input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.causal)

                attn_weights_0 = torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True)
                with torch.no_grad():
                    attn_weights_1 = attn_weights_0.clone().to(dtype).requires_grad_(True)
                total_mask = (~(
                    torch.tril(torch.randn((24, 24), device="cuda")).bool()
                ).unsqueeze(0).unsqueeze(0))

                with torch.cuda.amp.autocast(dtype=dtype):
                    actual = fused_fn(attn_weights_0, total_mask)
                    self.assertEqual(actual.dtype, dtype)
                expected = torch_fn(attn_weights_1, total_mask)
                torch.testing.assert_close(actual, expected)

                g0 = torch.randn_like(actual)
                with torch.no_grad():
                    g1 = g0.clone()
                actual.backward(g0)
                expected.backward(g1)
//  [python]
unary_ops = [
    hardswish,
    torch._C._nn.hardswish,
    torch.sigmoid,
    torch.reciprocal,
    torch.neg,
    torch.relu,
    torch.isnan,
    torch.log,
    torch.log10,
    torch.log1p,
    torch.log2,
    torch.exp,
    torch.expm1,
    torch.erf,
    torch.erfc,
    torch.cos,
    torch.sin,
    torch.tan,
    torch.acos,
    torch.asin,
    torch.cosh,
    torch.sinh,
    torch.atan,
    torch.tanh,
    torch.sqrt,
    torch.rsqrt,
    torch.abs,
    torch.ceil,
    torch.floor,
    torch.round,
    torch.trunc,
    torch.lgamma,
]

print("{:20s} {:>10s} {:>10s} {:>10s}".format("op", "eager", "nnc", "speedup"))

for op in unary_ops:
    x = torch.rand((1024, 1024))
    traced = torch.jit.trace(lambda x: op(x), (x))

    # Warmup.
    warmup_iters = 8
    for _ in range(warmup_iters):
        op(x)
        traced(x)

    # Validate result.
    torch.testing.assert_allclose(op(x), traced(x))

    # Benchmark.
    bench_iters = 100
    teager = timeit.timeit(stmt="op(x)", globals=globals(), number=bench_iters)
    tjit = timeit.timeit(stmt="traced(x)", globals=globals(), number=bench_iters)
    print(f"{op.__name__:20s} {teager:10.3f} {tjit:10.3f} {teager/tjit:10.2f}")
//  [python]
unary_ops = [
    hardswish,
    torch._C._nn.hardswish,
    torch.sigmoid,
    torch.reciprocal,
    torch.neg,
    torch.relu,
    torch.isnan,
    torch.log,
    torch.log10,
    torch.log1p,
    torch.log2,
    torch.exp,
    torch.expm1,
    torch.erf,
    torch.erfc,
    torch.cos,
    torch.sin,
    torch.tan,
    torch.acos,
    torch.asin,
    torch.cosh,
    torch.sinh,
    torch.atan,
    torch.tanh,
    torch.sqrt,
    torch.rsqrt,
    torch.abs,
    torch.ceil,
    torch.floor,
    torch.round,
    torch.trunc,
    torch.lgamma,
]

print("{:20s} {:>10s} {:>10s} {:>10s}".format("op", "eager", "nnc", "speedup"))

for op in unary_ops:
    x = torch.rand((1024, 1024))
    traced = torch.jit.trace(lambda x: op(x), (x))

    # Warmup.
    warmup_iters = 8
    for _ in range(warmup_iters):
        op(x)
        traced(x)

    # Validate result.
    torch.testing.assert_close(op(x), traced(x))

    # Benchmark.
    bench_iters = 100
    teager = timeit.timeit(stmt="op(x)", globals=globals(), number=bench_iters)
    tjit = timeit.timeit(stmt="traced(x)", globals=globals(), number=bench_iters)
    print(f"{op.__name__:20s} {teager:10.3f} {tjit:10.3f} {teager/tjit:10.2f}")
//  [python]
def test_fused_scale_mask_softmax(self):
        """
        attention_scores.shape = [4, 12, 24, 24]
        mask.shape = [4, 1, 24, 24]
        """
        for (dtype, scale, softmax_in_fp32) in itertools.product(
                (torch.half, torch.bfloat16),
                (None, 2.0),
                (False, True),
        ):
            with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"):
                input_in_fp16 = dtype == torch.half
                input_in_bf16 = dtype == torch.bfloat16
                if not (scale is None or softmax_in_fp32):
                    with self.assertRaises(RuntimeError):
                        self._setup_fused_softmax(input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.padding)
                    return
                fused_fn, torch_fn = self._setup_fused_softmax(input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.padding)

                attention_scores_0 = torch.randn((4, 12, 24, 24)).to(device="cuda", dtype=dtype).requires_grad_(True)
                with torch.no_grad():
                    attention_scores_1 = attention_scores_0.clone().requires_grad_(True)
                mask = torch.randint(0, 2, (4, 1, 24, 24), device="cuda").bool()
                expected = fused_fn(attention_scores_0, mask)
                actual = torch_fn(attention_scores_1, mask)
                torch.testing.assert_allclose(actual, expected)

                g0 = torch.rand_like(actual)
                with torch.no_grad():
                    g1 = g0.clone()
                expected.backward(g0)
                actual.backward(g1)
//  [python]
def test_fused_scale_mask_softmax(self):
        """
        attention_scores.shape = [4, 12, 24, 24]
        mask.shape = [4, 1, 24, 24]
        """
        for (dtype, scale, softmax_in_fp32) in itertools.product(
                (torch.half, torch.bfloat16),
                (None, 2.0),
                (False, True),
        ):
            with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"):
                input_in_fp16 = dtype == torch.half
                input_in_bf16 = dtype == torch.bfloat16
                if not (scale is None or softmax_in_fp32):
                    with self.assertRaises(RuntimeError):
                        self._setup_fused_softmax(input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.padding)
                    return
                fused_fn, torch_fn = self._setup_fused_softmax(input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.padding)

                attention_scores_0 = torch.randn((4, 12, 24, 24)).to(device="cuda", dtype=dtype).requires_grad_(True)
                with torch.no_grad():
                    attention_scores_1 = attention_scores_0.clone().requires_grad_(True)
                mask = torch.randint(0, 2, (4, 1, 24, 24), device="cuda").bool()
                expected = fused_fn(attention_scores_0, mask)
                actual = torch_fn(attention_scores_1, mask)
                torch.testing.assert_close(actual, expected)

                g0 = torch.rand_like(actual)
                with torch.no_grad():
                    g1 = g0.clone()
                expected.backward(g0)
                actual.backward(g1)
//  [python]
def test_bias_layernorm():
    class Projection(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(1024, 1024)
            self.dropout = nn.Dropout(0.2)
            self.ln = nn.LayerNorm(1024)

        def forward(self, x, residual):
            x = self.linear(x)
            x = self.dropout(x)
            x = self.ln(x + residual)
            return x

    mod = Projection().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))
    torch.testing.assert_allclose(sch.mod.linear.bias, mod.linear.bias)

    sch["linear"].decompose()
    sch.trace(flatten=True)
    print(sch.mod)

    def pattern(x, bias, residual):
        return F.layer_norm(F.dropout(x + bias) + residual, 1024)

    subgraph = sch.find(pattern)
    assert len(subgraph) == 1
    assert len(subgraph[0]) == 4
    assert subgraph[0][0][1].target == operator.add
    assert subgraph[0][1][1].target == "dropout"
    assert subgraph[0][2][1].target == operator.add
    assert subgraph[0][3][1].target == "ln"

    sch.fuse(subgraph, compiler="TorchScript", name="FusedLN")
    assert isinstance(sch["FusedLN_0"].mod, torch.jit.ScriptModule)

    sch_model, _ = slapo.build(sch, init_weights=False)
    torch.testing.assert_allclose(sch_model.linear.weight, mod.linear.weight)
    print(sch_model)

    inp = torch.randn((1, 16, 1024, 1024), requires_grad=True).cuda()
    resid = torch.randn((1, 16, 1024, 1024), requires_grad=True).cuda()
    reset_random_seeds()
    out = sch_model(inp, resid)

    reset_random_seeds()
    out_ref = mod(inp, resid)
    torch.testing.assert_close(out, out_ref)
//  [python]
def test_bias_layernorm():
    class Projection(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(1024, 1024)
            self.dropout = nn.Dropout(0.2)
            self.ln = nn.LayerNorm(1024)

        def forward(self, x, residual):
            x = self.linear(x)
            x = self.dropout(x)
            x = self.ln(x + residual)
            return x

    mod = Projection().cuda()
    sch = slapo.create_schedule(copy.deepcopy(mod))
    torch.testing.assert_close(sch.mod.linear.bias, mod.linear.bias)

    sch["linear"].decompose()
    sch.trace(flatten=True)
    print(sch.mod)

    def pattern(x, bias, residual):
        return F.layer_norm(F.dropout(x + bias) + residual, 1024)

    subgraph = sch.find(pattern)
    assert len(subgraph) == 1
    assert len(subgraph[0]) == 4
    assert subgraph[0][0][1].target == operator.add
    assert subgraph[0][1][1].target == "dropout"
    assert subgraph[0][2][1].target == operator.add
    assert subgraph[0][3][1].target == "ln"

    sch.fuse(subgraph, compiler="TorchScript", name="FusedLN")
    assert isinstance(sch["FusedLN_0"].mod, torch.jit.ScriptModule)

    sch_model, _ = slapo.build(sch, init_weights=False)
    torch.testing.assert_close(sch_model.linear.weight, mod.linear.weight)
    print(sch_model)

    inp = torch.randn((1, 16, 1024, 1024), requires_grad=True).cuda()
    resid = torch.randn((1, 16, 1024, 1024), requires_grad=True).cuda()
    reset_random_seeds()
    out = sch_model(inp, resid)

    reset_random_seeds()
    out_ref = mod(inp, resid)
    torch.testing.assert_close(out, out_ref)
//  [python]
def test_init_from_module(self):
        filter = ImageFilter(TestModule())
        out = filter(torch.zeros(1, 3, 3, 3))
        torch.testing.assert_allclose(torch.ones(1, 3, 3, 3), out)
//  [python]
def test_init_from_module(self):
        filter = ImageFilter(TestModule())
        out = filter(torch.zeros(1, 3, 3, 3))
        torch.testing.assert_close(torch.ones(1, 3, 3, 3), out)
//  [python]
def test_instances_in_sync_after_training():
    NUM_TOTAL_REPLICAS = 4
    NUM_INSTANCES = 2

    script_dir = Path(os.path.realpath(__file__)).parent.parent
    cmd = [
        "poprun",
        "--mpi-global-args",
        "--tag-output",
        "--num-replicas",
        str(NUM_TOTAL_REPLICAS),
        "--num-instances",
        str(NUM_INSTANCES),
        sys.executable,
        str(script_dir / "popdist_training.py"),
    ]

    with tempfile.TemporaryDirectory() as temp_dir:
        testing_util.run_command(cmd, cwd=temp_dir)

        checkpoint_0 = torch.load(
            os.path.join(temp_dir, "checkpoint-instance-0.pt"),
        )
        checkpoint_1 = torch.load(
            os.path.join(temp_dir, "checkpoint-instance-1.pt"),
        )

        # The final parameters should be the same for all instances.
        for k, v in checkpoint_0.items():
            if not (k.endswith("weight") or k.endswith("bias")):
                continue
            torch.testing.assert_allclose(v.data, checkpoint_1[k].data)
//  [python]
def test_instances_in_sync_after_training():
    NUM_TOTAL_REPLICAS = 4
    NUM_INSTANCES = 2

    script_dir = Path(os.path.realpath(__file__)).parent.parent
    cmd = [
        "poprun",
        "--mpi-global-args",
        "--tag-output",
        "--num-replicas",
        str(NUM_TOTAL_REPLICAS),
        "--num-instances",
        str(NUM_INSTANCES),
        sys.executable,
        str(script_dir / "popdist_training.py"),
    ]

    with tempfile.TemporaryDirectory() as temp_dir:
        testing_util.run_command(cmd, cwd=temp_dir)

        checkpoint_0 = torch.load(
            os.path.join(temp_dir, "checkpoint-instance-0.pt"),
        )
        checkpoint_1 = torch.load(
            os.path.join(temp_dir, "checkpoint-instance-1.pt"),
        )

        # The final parameters should be the same for all instances.
        for k, v in checkpoint_0.items():
            if not (k.endswith("weight") or k.endswith("bias")):
                continue
            torch.testing.assert_close(v.data, checkpoint_1[k].data)
//  [python]
def test_2d_filter_correctness(self, filter_name):
        "Test correctness of filters (2d only)"
        tfm = ImageFilter(filter_name, 3, **ADDITIONAL_ARGUMENTS)
        filter = tfm._get_filter_from_string(filter_name, size=3, ndim=2).filter.squeeze()
        torch.testing.assert_allclose(filter, EXPECTED_FILTERS[filter_name])
//  [python]
def test_2d_filter_correctness(self, filter_name):
        "Test correctness of filters (2d only)"
        tfm = ImageFilter(filter_name, 3, **ADDITIONAL_ARGUMENTS)
        filter = tfm._get_filter_from_string(filter_name, size=3, ndim=2).filter.squeeze()
        torch.testing.assert_close(filter, EXPECTED_FILTERS[filter_name])
//  [python]
def test_cornerfromparams(self):
        box = box_corners_from_param(torch.tensor([1, 2, 3]).float(), np.pi / 2, torch.tensor([1, 1, 1]))
        torch.testing.assert_allclose(
            box,
            1
            + torch.tensor(
                [
                    [1.0, -0.5, -1.5],
                    [1.0, 0.5, -1.5],
                    [-1.0, 0.5, -1.5],
                    [-1.0, -0.5, -1.5],
                    [1.0, -0.5, 1.5],
                    [1.0, 0.5, 1.5],
                    [-1.0, 0.5, 1.5],
                    [-1.0, -0.5, 1.5],
                ]
            ),
        )
//  [python]
def test_cornerfromparams(self):
        box = box_corners_from_param(torch.tensor([1, 2, 3]).float(), np.pi / 2, torch.tensor([1, 1, 1]))
        torch.testing.assert_close(
            box,
            1
            + torch.tensor(
                [
                    [1.0, -0.5, -1.5],
                    [1.0, 0.5, -1.5],
                    [-1.0, 0.5, -1.5],
                    [-1.0, -0.5, -1.5],
                    [1.0, -0.5, 1.5],
                    [1.0, 0.5, 1.5],
                    [-1.0, 0.5, 1.5],
                    [-1.0, -0.5, 1.5],
                ]
            ),
        )
//  [python]
def test_call_2d_prob_0(self, filter_name):
        "Text function `__call__` for 2d images"
        filter = RandImageFilter(filter_name, 3, 0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_IMAGE_2D)
        torch.testing.assert_allclose(out_tensor, SAMPLE_IMAGE_2D)
//  [python]
def test_call_2d_prob_0(self, filter_name):
        "Text function `__call__` for 2d images"
        filter = RandImageFilter(filter_name, 3, 0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_IMAGE_2D)
        torch.testing.assert_close(out_tensor, SAMPLE_IMAGE_2D)
//  [python]
def test_call_3d_prob_0(self, filter_name):
        "Text function `__call__` for 3d images"
        filter = RandImageFilter(filter_name, 3, 0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_IMAGE_3D)
        torch.testing.assert_allclose(out_tensor, SAMPLE_IMAGE_3D)
//  [python]
def test_call_3d_prob_0(self, filter_name):
        "Text function `__call__` for 3d images"
        filter = RandImageFilter(filter_name, 3, 0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_IMAGE_3D)
        torch.testing.assert_close(out_tensor, SAMPLE_IMAGE_3D)
//  [python]
def test_call_2d_prob_0(self, filter_name):
        filter = RandImageFilterd("image_2d", filter_name, 3, 0.0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_DICT)
        torch.testing.assert_allclose(out_tensor["image_2d"].shape[1:], SAMPLE_IMAGE_2D.shape[1:])
//  [python]
def test_call_2d_prob_0(self, filter_name):
        filter = RandImageFilterd("image_2d", filter_name, 3, 0.0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_DICT)
        torch.testing.assert_close(out_tensor["image_2d"].shape[1:], SAMPLE_IMAGE_2D.shape[1:])
//  [python]
def test_call_3d_prob_0(self, filter_name):
        filter = RandImageFilterd("image_3d", filter_name, 3, 0.0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_DICT)
        torch.testing.assert_allclose(out_tensor["image_3d"].shape[1:], SAMPLE_IMAGE_3D.shape[1:])
//  [python]
def test_call_3d_prob_0(self, filter_name):
        filter = RandImageFilterd("image_3d", filter_name, 3, 0.0, **ADDITIONAL_ARGUMENTS)
        out_tensor = filter(SAMPLE_DICT)
        torch.testing.assert_close(out_tensor["image_3d"].shape[1:], SAMPLE_IMAGE_3D.shape[1:])