Function test_batch_formatting_dynamic

dali/test/python/test_tensor_formatting.py:221–257

Test batch formatting for Dynamic API.

(test_case, device)

@params(*product(BATCH_TEST_CASES, DEVICES))
def test_batch_formatting_dynamic(test_case, device):
    """Test batch formatting for Dynamic API."""
    layout = test_case["layout"]
    num_samples = test_case["num_samples"]
    shapes = test_case["shapes"]
    should_summarize = test_case["should_summarize"]
    arrays = test_case["arrays"]

    if num_samples == 0:
        import nvidia.dali.types as types

        b = ndd.batch([], dtype=types.FLOAT, device=device)
    else:
        b = ndd.batch(arrays, layout=layout, device=device)

    for method_name, method in [("str", str), ("repr", repr)]:
        result = method(b)
        assert "Batch" in result, f'expected "Batch" in "{result}"'
        assert "dtype=" in result, f'expected "dtype=" in "{result}"'
        assert f'device="{device}"' in result, f'expected "device=\\"{device}\\"" in "{result}"'
        assert (
            f"num_samples={num_samples}" in result
        ), f'expected "num_samples={num_samples}" in "{result}"'
        if layout:
            assert f'layout="{layout}"' in result, f'expected "layout=\\"{layout}\\"" in "{result}"'

        if num_samples > 0:
            assert result.count("[") > 1, f'expected multiple "[" in "{result}"'
            assert (
                shapes[0] in result and shapes[-1] in result
            ), f'expected "{shapes[0]}" and "{shapes[-1]}" in "{result}"'
            expected_vals = _format_values(arrays[0]) + _format_values(arrays[-1])
            assert all(
                v in result for v in expected_vals
            ), f'expected values {expected_vals} in "{result}"'
        if should_summarize:
            assert "..." in result, f'expected "..." in "{result}"'