Methods of `ParallelBeamModel` can't be jitted

[bwohlberg]

Given their potential use in iterative reconstructions (I know mbirjax provides its own reconstruction algorithms, but there are valid use cases), it would be desirable to be able to jit methods forward_project and back_project, but this currently fails with a ConcretizationTypeError.

Minimal example:

import jax
import jax.numpy as jnp
import mbirjax

output_shape = (64, 256, 256)
angles = jnp.linspace(0, jnp.pi, output_shape[0])
model = mbirjax.ParallelBeamModel(output_shape, angles)
input_shape = model.get_params('recon_shape')
x = jnp.ones(input_shape, dtype=jnp.float32)
y = model.forward_project(x)  # succeeds

model.forward_project = jax.jit(model.forward_project)
y = model.forward_project(x)  # fails

Presumably this also applies to ConeBeamModel and other children of TomographyModel.

[bwohlberg]

This issue is still present in mbirjax 0.6.16.