Push pull

[hsalehipour]

Contributing Guidelines

• I have read and understood the CONTRIBUTING.md guidelines

Description

I have intentionally made several commits in this PR so that various issues can be easily documented. The main objective of this PR was to allow push and pull as options for doing LBM step in the JAX backend. I realized the latst PR had introduced some bugs. One was failing all warp issues (see the fix in indices masker). All linting and ruff issues appeared afterwards when doing ruff format . and ruff check . . There were also couple of old pytest issues that I had to fix to pass all tests (unrelated to this PR). I made sure the out of core LDC example runs nicely on my desktop.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

How Has This Been Tested?

• All pytest tests pass

Linting and Code Formatting

Make sure the code follows the project's linting and formatting standards. This project uses Ruff for linting.

To run Ruff, execute the following command from the root of the repository:

ruff check .

• Ruff passes

[mehdiataei]

Thanks.

1. Can you combine cosmetic changes into a single commit?

2. Why do we need the push implementation? I also notice it's inly for jax. Is there a benefit here? Do we have this for warp?

I recall the push method is not be compatible with a lot of our BCs such as the way we handle memories for aux data. If you want to add this as a feature it's great, but it needs to be fully compatible with our BCs otherwise it will be used incorrectly, especially for a fundamental thing like this (especially for warp)

Maybe the right choice to remove any push implementation (eg from warp if we have it) instead until we add it fully.

[hsalehipour]

Thanks. I combined the cosmetic commits into one. The push method is not implemented for Warp. It is only implemented in JAX here to allow the stepper function to be complete (without an extra apply BC outside of stepper in the case of pull method). We don't have any BC auxiliary data handling in JAX. That is only in WARP and I don't have a push method for WARP backend in this PR, so there are no issues.