x_shape = (3, 3), y_shape = (3, 3)

    @pytest.mark.parametrize(
        ["x_shape", "y_shape"],
        [
            [(3, 3), (3, 3)],
            [(32, 32), (32, 32)],
            [(64, 64), (64, 64)],
        ],
    )
    def test_example_graph(x_shape: tuple, y_shape: tuple):
>       run_graph_test_with_random_inputs(example_graph, [x_shape, y_shape])

tests/jax/graphs/test_example_graph.py:27: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/infra/graph_tester.py:48: in run_graph_test_with_random_inputs
    tester.test_with_random_inputs(graph, input_shapes)
tests/infra/op_tester.py:45: in test_with_random_inputs
    self.test(workload)
tests/infra/op_tester.py:35: in test
    self._compare(tt_res, cpu_res)
tests/infra/base_tester.py:55: in _compare
    compare_atol(device_output, golden_output, self._comparison_config.atol)
tests/infra/device_runner.py:147: in wrapper
    return DeviceRunner.run_on_cpu(workload)
tests/infra/device_runner.py:26: in run_on_cpu
    return DeviceRunner._run_on_device(DeviceType.CPU, workload)
tests/infra/device_runner.py:72: in _run_on_device
    return device_workload.execute()
tests/infra/utils.py:27: in execute
    return self.executable(*self.args, **self.kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

device_output = Array([[1.0001237, 1.0001053, 1.0001491],
       [1.0001608, 1.0001174, 1.0001533],
       [1.0000875, 1.0000967, 1.0001339]], dtype=float32)
golden_output = Array([[1.6487212, 1.6487212, 1.6487212],
       [1.6487212, 1.6487212, 1.6487212],
       [1.6487212, 1.6487212, 1.6487212]], dtype=float32)
atol_config = AtolConfig(enabled=True, required_atol=0.1)

    @run_on_cpu
    def compare_atol(
        device_output: Tensor, golden_output: Tensor, atol_config: AtolConfig
    ) -> None:
        assert isinstance(device_output, jax.Array) and isinstance(
            golden_output, jax.Array
        ), f"Currently only jax.Array is supported {type(device_output)}, {type(golden_output)}"
    
        atol = jnp.max(jnp.abs(device_output - golden_output))
    
        assert (
>           atol <= atol_config.required_atol
        ), f"Atol comparison failed. Calculated atol={atol}"
E       AssertionError: Atol comparison failed. Calculated atol=0.6486337184906006

tests/infra/comparison.py:94: AssertionError

x_shape = (64, 64), y_shape = (64, 64)

    @pytest.mark.parametrize(
        ["x_shape", "y_shape"],
        [
            [(3, 3), (3, 3)],
            [(32, 32), (32, 32)],
            [(64, 64), (64, 64)],
        ],
    )
    def test_example_graph(x_shape: tuple, y_shape: tuple):
>       run_graph_test_with_random_inputs(example_graph, [x_shape, y_shape])

tests/jax/graphs/test_example_graph.py:27: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/infra/graph_tester.py:48: in run_graph_test_with_random_inputs
    tester.test_with_random_inputs(graph, input_shapes)
tests/infra/op_tester.py:45: in test_with_random_inputs
    self.test(workload)
tests/infra/op_tester.py:35: in test
    self._compare(tt_res, cpu_res)
tests/infra/base_tester.py:55: in _compare
    compare_atol(device_output, golden_output, self._comparison_config.atol)
tests/infra/device_runner.py:147: in wrapper
    return DeviceRunner.run_on_cpu(workload)
tests/infra/device_runner.py:26: in run_on_cpu
    return DeviceRunner._run_on_device(DeviceType.CPU, workload)
tests/infra/device_runner.py:72: in _run_on_device
    return device_workload.execute()
tests/infra/utils.py:27: in execute
    return self.executable(*self.args, **self.kwargs)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

device_output = Array([[1.6390116, 1.6390126, 1.6390126, ..., 1.0001031, 1.0001142,
        1.0001161],
       [1.6390125, 1.6390126, ...    1.0001127],
       [1.0000914, 1.000103 , 1.0001153, ..., 1.0001247, 1.0001589,
        1.0000964]], dtype=float32)
golden_output = Array([[1.6487212, 1.6487212, 1.6487212, ..., 1.6487212, 1.6487212,
        1.6487212],
       [1.6487212, 1.6487212, ...    1.6487212],
       [1.6487212, 1.6487212, 1.6487212, ..., 1.6487212, 1.6487212,
        1.6487212]], dtype=float32)
atol_config = AtolConfig(enabled=True, required_atol=0.1)

    @run_on_cpu
    def compare_atol(
        device_output: Tensor, golden_output: Tensor, atol_config: AtolConfig
    ) -> None:
        assert isinstance(device_output, jax.Array) and isinstance(
            golden_output, jax.Array
        ), f"Currently only jax.Array is supported {type(device_output)}, {type(golden_output)}"
    
        atol = jnp.max(jnp.abs(device_output - golden_output))
    
        assert (
>           atol <= atol_config.required_atol
        ), f"Atol comparison failed. Calculated atol={atol}"
E       AssertionError: Atol comparison failed. Calculated atol=0.6486340761184692

tests/infra/comparison.py:94: AssertionError