Error initializing tfb.AutoregressiveNetwork using jax substrate

Trying to initialize an instance of tfb.AutoregressiveNetwork using the jax substrate fails with an AttributeError.

With the example usage from the docs:

from tensorflow_probability.substrates import jax as tfp
tfb = tfp.bijector

tfb.AutoregressiveNetwork(params=2, hidden_units=[10,10])

raises the error:

File .../site-packages/tensorflow_probability/substrates/jax/bijectors/masked_autoregressive.py:967, in AutoregressiveNetwork.__init__(self, params, event_shape, conditional, conditional_event_shape, conditional_input_layers, hidden_units, input_order, hidden_degrees, activation, use_bias, kernel_initializer, bias_initializer, kernel_regularizer, bias_regularizer, kernel_constraint, bias_constraint, validate_args, **kwargs)
    965 self._kernel_regularizer = kernel_regularizer
    966 self._bias_regularizer = bias_regularizer
--> 967 self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
    968 self._bias_constraint = bias_constraint
    969 self._validate_args = validate_args

AttributeError: module 'tensorflow_probability.python.internal.backend.jax.numpy_keras' has no attribute 'constraints'

Tested with:
tfp version: 0.19.0
jax version: 0.4.4 and 0.3.25

Start with tensorflow_probability/substrates/jax/bijectors/masked_autoregressive.py around the AutoregressiveNetwork initializer and reproduce the documented example. Check how the JAX backend handles kernel constraints; done means the example initializes without the reported AttributeError and has a regression test for the behavior.