Give compile_training_step's shared parameters a static shape for JAX and MLX

_make_shared_params creates the trainable slots with pytensor.shared(init_val, name=...) and no shape, so an RV declared with dims and no shape stays (None, None) after replacement. That is fine on the C and Numba backends. Under JAX and MLX an arange whose bound comes from such a parameter's shape is not concrete and the dispatch raises, which is what pt.trace over W @ W.T in a categorical trace penalty lowers to in the SVGP gradient. Declaring shape= alongside dims= on the model side avoids it, so this is a convenience for models that did not.

import numpy as np
import pymc as pm
import pytensor.tensor as pt
import ptgp as pg

rng = np.random.default_rng(0)
X = np.column_stack([rng.uniform(0, 10, 200), rng.integers(0, 4, 200)]).astype(np.float64)
y = rng.normal(size=200)
X_var = pt.matrix("X", shape=(None, 2))
y_var = pt.vector("y")

with pm.Model(coords={"level": list("abcd"), "rank": [0, 1]}) as model:
    W = pm.Normal("W", dims=("level", "rank"))  # dims only, so W's type shape is (None, None)
    kappa = pm.LogNormal("kappa", dims="level")
    pm.Potential("trace_pen", -pt.square(pt.log(pt.trace(W @ W.T + pt.diag(kappa)))))
    kernel = pg.kernels.ExpQuad(input_dim=2, ls=1.0, active_dims=[0]) * pg.kernels.LowRankCategorical(
        input_dim=2, num_levels=4, W=W, kappa=kappa, active_dims=[1]
    )
    gp = pg.gp.SVGP(
        kernel=kernel,
        likelihood=pg.likelihoods.Gaussian(1.0),
        inducing_variable=pg.inducing.Points(pt.as_tensor_variable(X[:8])),
        variational_params=pg.gp.init_variational_params(8),
    )

train_step, *_ = pg.optim.compile_training_step(
    lambda gp_, X_, y_: pg.objectives.elbo(gp_, X_, y_, n_data=200),
    gp,
    X_var,
    y_var,
    model=model,
    compile_kwargs={"mode": "JAX"},
)  # NotImplementedError: JAX requires the arguments of `jax.numpy.arange` to be constants.
# Compiles with mode="NUMBA", or with W declared as
# pm.Normal("W", dims=("level", "rank"), shape=(4, 2)).

The init arrays are the parameters, and parameters do not resize during a fit, so the shape is always known at that point:

# ptgp/optim/training.py, _make_shared_params
shared_params[vv] = pytensor.shared(init_val, name=vv.name, shape=init_val.shape)
...
init_array = np.asarray(init, dtype=np.float64)
shared_extras.append(pytensor.shared(init_array, name=var.name, shape=init_array.shape))

Checked on ptgp 0.1.2 with pytensor 3.2.3, pymc 6.2.0, jax 0.9.0.1. pytensor main has the same ARange dispatch, and ptgp 0.1.2 pins pytensor<3.3, so newer pytensor was not tried.

Start in ptgp/optim/training.py at _make_shared_params and follow compile_training_step's shared-parameter setup. Run the reproducer with dims-only parameters under JAX and MLX, then verify that parameters and shared extras retain their initialization shapes and that the JAX compilation succeeds without changing the existing NUMBA behavior.