def test_to_device():
        cpu_array = random_input_tensor((32, 32))
        device = jax.devices()[0]
        tt_array = jax.device_put(cpu_array, device)
>       assert tt_array.device.device_kind == "wormhole"
E       AssertionError: assert 'cpu' == 'wormhole'
E         
E         - wormhole
E         + cpu

tests/TTIR/test_device.py:21: AssertionError

comparison_config = ComparisonConfig(equal=EqualConfig(enabled=False), atol=AtolConfig(enabled=False, required_atol=0.1), pcc=PccConfig(enabled=True, required_pcc=0.99), allclose=AllcloseConfig(enabled=False, rtol=0.01, atol=0.01))

    @pytest.fixture
    def inference_tester(
        comparison_config: ComparisonConfig,
    ) -> FlaxDistilBertForMaskedLMTester:
>       return FlaxDistilBertForMaskedLMTester(comparison_config)
E       TypeError: Can't instantiate abstract class FlaxDistilBertForMaskedLMTester with abstract methods _get_activations, _get_forward_method_args, _get_forward_method_kwargs, _get_forward_method_name

tests/jax/models/flax_distil_bert_for_masked_lm/test_flax_distil_bert_for_masked_lm.py:53: TypeError

jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

inference_tester = <test_simple_nn.SimpleNNTester object at 0x7f061c367640>

    def test_simple_nn_inference(inference_tester: SimpleNNTester):
>       inference_tester.test()

tests/jax/models/test_simple_nn.py:131: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/infra/model_tester.py:122: in test
    self._test_inference()
tests/infra/model_tester.py:138: in _test_inference
    tt_res = DeviceRunner.run_on_tt_device(compiled_workload)
tests/infra/device_runner.py:21: in run_on_tt_device
    return DeviceRunner._run_on_device(DeviceType.TT, workload)
tests/infra/device_runner.py:73: in _run_on_device
    return device_workload.execute()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = Workload(executable=<PjitFunction of <bound method SimpleNN.__call__ of SimpleNN(
  w0=Array(shape=(784, 128), dtype=f...0., 0., 0., 0., 0., 0.,
        0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],      dtype=float32)})

    def execute(self) -> Any:
>       return self.executable(*self.args, **self.kwargs)
E       TypeError: SimpleNN.__call__() got multiple values for argument 'act'

tests/infra/utils.py:26: TypeError