feature/jax_cpu_batch_size_1

[Jammy2211]

This pull request introduces a minor update to the initialization logic in autofit/non_linear/fitness.py, specifically to improve performance when using the JAX backend on CPUs. Now, when JAX is detected as the analysis backend and it is running on a CPU, the batch size is automatically set to 1, and a log message is generated to inform the user.

This is because CPU's do not get any speed up from using a vmap with a batch size above 1, but this can cause to issues with shared parallelism.

Performance optimization for JAX CPU backend:

• Automatically sets batch_size to 1 and logs an informational message when JAX is used with the CPU backend to improve performance. (autofit/non_linear/fitness.py)

Minor cleanup:

• Removes the unused xp parameter from the __init__ method signature. (autofit/non_linear/fitness.py)

[Copilot]

Pull request overview

Updates Fitness initialization to better handle JAX-on-CPU scenarios by forcing a batch size of 1 and removing an xp initializer argument.

Changes:

• Removed the xp parameter from Fitness.__init__.
• Added JAX CPU-backend detection that logs an info message and sets batch_size = 1 when analysis._use_jax is enabled.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Setting the local batch_size = 1 here only changes self.batch_size and does not actually constrain the size of batches passed into the vmapped likelihood (the effective vmap batch size is determined by the shape of the parameters array passed to call_wrap). If the goal is to avoid vmapping over batches >1 on CPU, this needs to be enforced at the call site (e.g., configuring vectorized samplers to use n_batch=1) or by adding logic in call_wrap to chunk/split incoming 2D parameter arrays to the desired batch size.

[Copilot]

This introduces new CPU-specific behavior for JAX (forcing batch_size to 1) but there doesn’t appear to be coverage asserting this branch (or that the informational log is emitted). Consider adding a unit test that simulates a JAX CPU backend (e.g., via monkeypatching jax.default_backend or a stub module) and verifies the resulting Fitness.batch_size value.

[Copilot]

The xp parameter was removed from __init__, but xp is still referenced later in this method (e.g., -xp.inf for resample_figure_of_merit / quick_update_max_lh). This will raise NameError at runtime during Fitness construction. Replace remaining xp references with np or self._xp (consistent with the existing _xp property) and ensure resample_figure_of_merit / quick_update_max_lh are initialized using that backend.