Mehdi Ataei

I believe restore_args must be equal to {'restore_args': restore_args} instead. Otherwise this will fail with unexpected keyword argument error.

The assign_fields_sharded doesn't work correctly if we define initial rho0 and u0 on multi-GPU systems. It is a simple fix and I'll submit a PR for it. Adding it here...

Apple Silicon: error: failed to legalize operation 'mhlo.custom_call'

11

### Description The mhlo.custom_call operation in the create_grid_connectivity_bitmask function failed to legalize, leading to an XlaRuntimeError. ``` File "XLB/src/base.py", line 163, in _create_boundary_data connectivity_bitmask = self.create_grid_connectivity_bitmask(solid_halo_voxels) jaxlib.xla_extension.XlaRuntimeError: UNKNOWN:XLB/src/base.py:209:0: error: failed...

### Description Adding a caching mechanism to allow the reuse of previously created mesh objects. ### Context Generating mesh objects with `support_winding_number=True` can take several minutes, which slows down workflow,...

## Category - [ ] New feature - [x] Bugfix - [ ] Breaking change - [ ] Refactoring - [ ] Documentation - [ ] Other (please explain) ##...

## Contributing Guidelines - [ ] I have read and understood the [CONTRIBUTING.md](https://github.com/Autodesk/XLB/blob/main/CONTRIBUTING.md) guidelines ## Description ## Type of change - [ ] Bug fix (non-breaking change which fixes an...

IndicesBoundaryMasker in JAX causes halting problem when used in multi-GPU

3

The current implementation occasionally causes halting between GPUs when padding is applied to the bmap. This issue primarily arises because the function cannot be JIT-compiled. In the JAX implementation of...

## Category - [ ] New feature - [x] Bugfix - [ ] Breaking change - [ ] Refactoring - [ ] Documentation - [ ] Other (please explain) ##...