Atlas helps TensorFlow developers in 2026 write unit tests for untested code by reading existing conventions, generating new test files, and running them with `pytest`, ensuring your `tf.function` and `tf.data` pipelines are robust and well-tested.
How does Atlas identify untested TensorFlow functions?
Atlas identifies untested TensorFlow functions by first reading the target module and then using its `lsp` tool to enumerate exported symbols. This process ensures that every public `tf.function` or Keras 3 model definition is considered for testing, preventing any critical component from being overlooked in 2026.
To begin adding tests, Atlas uses its `read` tool to ingest the contents of your TensorFlow module, such as `src/models/my_keras_model.py`. Following this, the `lsp` tool's `documentSymbol` operation is invoked. This operation leverages the Language Server Protocol to parse the Abstract Syntax Tree (AST) of your Python files, identifying all exported functions, classes, and methods. For a TensorFlow project, this means Atlas can precisely list every `tf.function` decorated method, custom Keras layer, or `tf.data` pipeline component that needs unit test coverage. This method, powered by tree-sitter, is far more accurate than simple line-based scanning, ensuring a comprehensive understanding of your module's public API, especially within a `pyproject.toml` defined project pinning `tensorflow` and `keras`.
How does Atlas match existing TensorFlow test conventions?
Atlas ensures new TensorFlow unit tests match your repository's existing conventions by using its `grep` tool to find and analyze existing test files. This approach, crucial for maintaining code consistency in 2026, allows Atlas to replicate your project's `pytest` setup, `tf.test.TestCase` usage, and import styles, rather than inventing new ones.
Maintaining a consistent testing style is vital for large TensorFlow codebases. Atlas achieves this by first using its `grep` tool to search your project's `tests/` directory for existing test files, typically named `test_*.py`. It looks for patterns like `import tensorflow as tf` and `class MyTest(tf.test.TestCase):` to understand your project's specific `pytest` configuration and testing idioms. Atlas then analyzes the structure, import statements, and assertion patterns within these files. For instance, if your project uses `tf.test.TestCase` for its test classes and a specific fixture setup, Atlas will replicate these conventions when drafting new test files like `tests/test_my_module.py`. This ensures that any new tests Atlas writes direct integrate with your existing suite, making them immediately recognizable and maintainable by other TensorFlow developers.
How does Atlas safely write new TensorFlow unit tests?
Atlas writes new TensorFlow unit test specifications safely using its `write` tool, which always presents a unified diff for approval before any changes are committed to disk. This permission-gated process, a core safety feature in 2026, ensures you review every line of the proposed `test_my_module.py` file, including `pytest` assertions and `tf.function` calls, before it becomes part of your codebase.
Once Atlas has identified the functions to test and understood your project's testing conventions, it drafts the new test file using the `write` tool. For example, if you're testing `src/data/input_pipeline.py`, Atlas might propose `tests/test_input_pipeline.py`. Before this file is written, Atlas computes a unified diff, clearly showing every addition, modification, or deletion. This diff is then presented to you in a permission prompt. You have the explicit choice to `allow`, `ask` for clarification, or `deny` the changes. This granular control is crucial for sensitive operations, allowing you to verify the correctness of `pytest` fixtures, `tf.data` pipeline assertions, and `tf.function` test cases. Atlas can also format the diff using `black` to ensure it adheres to your project's code style, further enhancing readability and reviewability.
How does Atlas run and iterate on TensorFlow tests?
Atlas runs new TensorFlow unit tests using its `bash` tool, executing `uv run pytest` to get immediate feedback on test failures. If the output exceeds 2000 lines or 50 KB, Atlas truncates it for readability but saves the full log to a file, allowing developers in 2026 to efficiently debug complex `tf.data` pipeline issues or `tf.function` retracing warnings.
The true value of a test lies in its execution. After writing the initial test spec, Atlas uses its `bash` tool to run the `pytest` suite. The command `uv run pytest tests/test_my_module.py` is executed directly within your project's environment, ensuring all dependencies are correctly managed by `uv`. Atlas captures the output, providing immediate feedback on test successes or failures. For extensive test runs, where output might exceed 2000 lines or 50 KB, Atlas intelligently truncates the display in the terminal for better readability, while simultaneously saving the complete log to a file for detailed analysis. If tests fail, Atlas facilitates an iterative process using the `edit` tool to modify the test file directly. For larger modules or complex `tf.data` pipelines requiring multiple fixes, Atlas can maintain a `todowrite` list, helping you track progress and ensure all identified issues are addressed until the entire suite is green.
How Atlas optimizes TensorFlow code during testing?
Atlas goes beyond just writing tests; it can also optimize your TensorFlow code, especially for performance-critical `tf.function` and `tf.data` pipelines. In 2026, Atlas can wrap hot Python loops in `tf.function` and add `tf.data` prefetch and cache stages, then measure the step time before and after, providing concrete performance improvements alongside robust testing.
Atlas's understanding of TensorFlow extends to its performance characteristics. When working with modules that define `tf.function` graph tracing or `tf.data` input pipelines, Atlas can identify opportunities for optimization. For instance, if a Python loop is identified as a bottleneck, Atlas can suggest and apply the `tf.function` decorator, explaining the retracing warnings it removes and the performance benefits gained. Similarly, for `tf.data` pipelines, Atlas can recommend and implement `prefetch` and `cache` stages to improve data loading efficiency. Crucially, Atlas can then measure the step time before and after these changes, providing empirical evidence of the optimization's impact. This capability ensures that your TensorFlow code is not only well-tested but also performs optimally, leveraging the full power of Keras 3 and TensorFlow's graph execution.
Step by step
- 01Use `atlas read src/my_module.py` to ingest the TensorFlow module requiring tests.
- 02Run `atlas lsp documentSymbol src/my_module.py` to enumerate all public `tf.function` and Keras model symbols.
- 03Execute `atlas grep "class .*\(tf.test.TestCase\)"` in your `tests/` directory to identify existing `pytest` and `tf.test.TestCase` conventions.
- 04Draft the new test file, e.g., `tests/test_my_module.py`, with `atlas write tests/test_my_module.py`, reviewing the unified diff for approval.
- 05Run the newly drafted tests using `atlas bash "uv run pytest tests/test_my_module.py"` to check for failures.
- 06If tests fail, use `atlas edit tests/test_my_module.py` to make necessary corrections to the test logic or assertions.
- 07Ensure code style consistency by running `atlas bash "uv run black tests/test_my_module.py"` and reviewing the diff.
- 08For large modules, use `atlas todowrite "Address remaining test failures in my_module"` to track progress on complex `tf.data` or `tf.function` test cases.
Frequently asked questions
- How does Atlas handle `tf.function` retracing warnings during testing?
- Atlas can identify hot Python loops and suggest wrapping them in `tf.function`, explaining the retracing warnings it removes. It helps ensure your `tf.function` decorated code is tested correctly and efficiently.
- Can Atlas help test `tf.data` pipelines effectively?
- Yes, Atlas understands `tf.data` pipelines. It can help write tests for them, add `prefetch` and `cache` stages for optimization, and measure step times to validate performance improvements.
- What test runner does Atlas use for TensorFlow projects?
- Atlas uses `pytest` as the test runner for TensorFlow projects. It executes `pytest` commands via its `bash` tool, ensuring compatibility with standard Python testing practices.
- How does Atlas ensure new tests match my TensorFlow repo's style?
- Atlas uses its `grep` tool to analyze existing `test_*.py` files, learning your project's `tf.test.TestCase` usage, import styles, and naming conventions. It then applies these learned patterns to new test files.
- Is Atlas safe to use with my production TensorFlow codebase?
- Yes, Atlas is designed with safety in mind. Every tool call is permission-gated, and all file edits are presented as a unified diff for your explicit approval before anything is written to disk.
- Can Atlas help optimize my Keras 3 model tests?
- Atlas can read your Keras 3 model definitions and `tf.data` input pipelines. It can assist in writing tests for custom layers and training steps, and suggest `tf.function` optimizations for performance.
- What package manager does Atlas use for TensorFlow dependencies?
- Atlas leverages `uv` as the package manager for TensorFlow projects. It ensures that `pytest` and other dependencies are correctly installed and managed within your project's environment.
Try Atlas in your terminal
The terminal-native AI coding agent. Free core, single binary.
Install AtlasRelated guides
Write Unit Tests for Untested Code with Atlas in 2026
How to write unit tests for untested code with Atlas in 2026: the lsp tool enumerates exported symbols, grep copies repo conventions, and bash actually runs the suite.
Atlas for TensorFlow: Terminal-Native AI Coding for Keras 3 and tf.data in 2026
Atlas is a terminal-native AI coding agent for TensorFlow and Keras 3 in 2026, where tf.function graph tracing and tf.data pipelines are what make training fast.
Onboard to an Unfamiliar Codebase in TensorFlow with Atlas in 2026
Quickly build a mental model of any TensorFlow repository in 2026 using Atlas. Leverage semantic search, safe exploration, and direct interaction with `pytest`, `uv`, and `black`.
Research a third-party API before integrating it in TensorFlow with Atlas in 2026
Streamline TensorFlow API integration in 2026. Atlas helps developers research external APIs, fetch documentation, and verify against `pyproject.toml` conventions before writing code.
Extract a shared helper from duplicated code in TensorFlow with Atlas in 2026
Refactor TensorFlow code in 2026: Use Atlas to find duplicated logic, extract it into a shared helper, and apply atomic patches. Ensure code quality with `pytest` and `black`.
Document a TensorFlow Module with a README in 2026
Atlas helps TensorFlow developers in 2026 generate accurate READMEs by reading live code, using `lsp` for APIs, and `bash` to verify samples. Ensure your TensorFlow documentation reflects current implementation, not
Run Atlas Headless in CI with TensorFlow in 2026
Automate TensorFlow model optimization and testing in CI pipelines using Atlas. Learn to run Atlas headless with `uv`, `pytest`, and `black` for machine-readable output in 2026.
Audit a TensorFlow Repository with Parallel Subagents in Atlas (2026)
Sweep your TensorFlow and Keras 3 codebase for issues without context window limits. Atlas uses parallel subagents and real tools like pytest and black for efficient, safe audits.