Enable jax.Hyperelastic(..., parallel=True) with jax.pmap on quadrature-points axis

[adtzlr]

This makes the Hyperelastic/Material JAX-classes with state variables a bit more unflexible: -1-batch axes are not allowed in pmap and hence, explicit axes (>0) must be specified. That means scalar-valued state variables must be converted to a length-1 array.

To use multiple cores of a CPU device as local XLA devices:

import os

os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.28%. Comparing base (4ee3496) to head (48dc44a).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #896   +/-   ##
=======================================
  Coverage   97.28%   97.28%           
=======================================
  Files         171      171           
  Lines        6149     6159   +10     
=======================================
+ Hits         5982     5992   +10     
  Misses        167      167           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.