[tx] Separate JaxBackend tests from TinkerEngine tests

[raulchen]

Summary

Separates backend-specific tests from engine-level tests, preparing for multiple backend implementations.

Changes

• Removed 7 tests that are specific to JaxBackend from test_engine.py to test_jax_backend.py
• Kept 2 tests that test engine-specific logic
• engine.py: extracted prepare_model_pass_batch() and prepare_sample_batch() to module level for test imports

Test coverage remains the same as before.
Note, there are still some functions not covered by current tests. But that's out of this PR's scope.

[gemini-code-assist]

Code Review

This pull request effectively separates backend-specific tests from engine-level tests, which is a great step towards supporting multiple backend implementations. The changes are well-executed: 7 Jax-specific tests have been moved from test_engine.py to test_jax_backend.py, and the remaining tests in test_engine.py are indeed engine-specific. The refactoring in engine.py to extract prepare_model_pass_batch() and prepare_sample_batch() into module-level functions is a clean way to enable their use in the new backend tests. The moved tests have been correctly adapted to interact with the JaxBackend directly. Overall, this is a solid refactoring that improves the test structure and maintainability. The code quality is high, and I have no further recommendations for changes.

[pcmoritz]

This fixes #819