Bug in tf-nightly/tfp-nightly: tensorflow_probability.python.internal.backend.jax.keras has no attribute 'constraints'

Example of how to reproduce the bug:

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

• [ x ] If applicable, include full error messages/tracebacks.

# altosaar @ host in ~/projects/variational-autoencoder on git:inverse-autoregressive-flow x [16:30:18] C:130
$ pip install --rupgradetf-nightly
Requirement already satisfied: tf-nightly in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (2.6.0.dev20210525)
Collecting tf-nightly
  Downloading tf_nightly-2.6.0.dev20210526-cp39-cp39-manylinux2010_x86_64.whl (450.8 MB)
     |████████████████████████████████| 450.8 MB 4.4 kB/s 
Requirement already satisfied: six~=1.15.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.15.0)
Requirement already satisfied: tb-nightly~=2.6.0.a in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (2.6.0a20210525)
Requirement already satisfied: protobuf>=3.9.2 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (3.17.1)
Requirement already satisfied: gast==0.4.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (0.4.0)
Requirement already satisfied: wheel~=0.35 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (0.36.2)
Requirement already satisfied: opt-einsum~=3.3.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (3.3.0)
Requirement already satisfied: tf-estimator-nightly~=2.5.0.dev in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (2.5.0.dev2021032601)
Requirement already satisfied: h5py~=3.1.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (3.1.0)
Requirement already satisfied: absl-py~=0.10 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (0.12.0)
Requirement already satisfied: grpcio<2.0,>=1.37.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.38.0)
Requirement already satisfied: typing-extensions~=3.7.4 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (3.7.4.3)
Requirement already satisfied: flatbuffers~=1.12.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.12)
Requirement already satisfied: keras-nightly~=2.6.0.dev in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (2.6.0.dev2021052500)
Requirement already satisfied: astunparse~=1.6.3 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.6.3)
Requirement already satisfied: wrapt~=1.12.1 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.12.1)
Requirement already satisfied: keras-preprocessing~=1.1.2 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.1.2)
Requirement already satisfied: google-pasta~=0.2 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (0.2.0)
Requirement already satisfied: termcolor~=1.1.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.1.0)
Requirement already satisfied: numpy~=1.19.2 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tf-nightly) (1.19.5)
Requirement already satisfied: werkzeug>=0.11.15 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (2.0.1)
Requirement already satisfied: markdown>=2.6.8 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (3.3.4)
Requirement already satisfied: google-auth<2,>=1.6.3 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (1.30.1)
Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (1.8.0)
Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (0.4.4)
Requirement already satisfied: requests<3,>=2.21.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (2.25.1)
Requirement already satisfied: setuptools>=41.0.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (52.0.0.post20210125)
Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tb-nightly~=2.6.0.a->tf-nightly) (0.6.1)
Requirement already satisfied: rsa<5,>=3.1.4 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from google-auth<2,>=1.6.3->tb-nightly~=2.6.0.a->tf-nightly) (4.7.2)
Requirement already satisfied: cachetools<5.0,>=2.0.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from google-auth<2,>=1.6.3->tb-nightly~=2.6.0.a->tf-nightly) (4.2.2)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from google-auth<2,>=1.6.3->tb-nightly~=2.6.0.a->tf-nightly) (0.2.8)
Requirement already satisfied: requests-oauthlib>=0.7.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from google-auth-oauthlib<0.5,>=0.4.1->tb-nightly~=2.6.0.a->tf-nightly) (1.3.0)
Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from pyasn1-modules>=0.2.1->google-auth<2,>=1.6.3->tb-nightly~=2.6.0.a->tf-nightly) (0.4.8)
Requirement already satisfied: chardet<5,>=3.0.2 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from requests<3,>=2.21.0->tb-nightly~=2.6.0.a->tf-nightly) (4.0.0)
Requirement already satisfied: idna<3,>=2.5 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from requests<3,>=2.21.0->tb-nightly~=2.6.0.a->tf-nightly) (2.10)
Requirement already satisfied: certifi>=2017.4.17 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from requests<3,>=2.21.0->tb-nightly~=2.6.0.a->tf-nightly) (2020.12.5)
Requirement already satisfied: urllib3<1.27,>=1.21.1 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from requests<3,>=2.21.0->tb-nightly~=2.6.0.a->tf-nightly) (1.26.4)
Requirement already satisfied: oauthlib>=3.0.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tb-nightly~=2.6.0.a->tf-nightly) (3.1.0)
Installing collected packages: tf-nightly
  Attempting uninstall: tf-nightly
    Found existing installation: tf-nightly 2.6.0.dev20210525
    Uninstalling tf-nightly-2.6.0.dev20210525:
      Successfully uninstalled tf-nightly-2.6.0.dev20210525
Successfully installed tf-nightly-2.6.0.dev20210526
(/scratch/gpfs/altosaar/environment-jax) 
# altosaar @ host in ~/projects/variational-autoencoder on git:inverse-autoregressive-flow x [16:31:59] 
$ pip install tfp-nightly --upgrade
Requirement already satisfied: tfp-nightly in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (0.14.0.dev20210525)
Collecting tfp-nightly
  Downloading tfp_nightly-0.14.0.dev20210527-py2.py3-none-any.whl (5.5 MB)
     |████████████████████████████████| 5.5 MB 18.9 MB/s 
Requirement already satisfied: cloudpickle>=1.3 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tfp-nightly) (1.6.0)
Requirement already satisfied: numpy>=1.13.3 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tfp-nightly) (1.19.5)
Requirement already satisfied: gast>=0.3.2 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tfp-nightly) (0.4.0)
Requirement already satisfied: six>=1.10.0 in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tfp-nightly) (1.15.0)
Requirement already satisfied: dm-tree in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tfp-nightly) (0.1.6)
Requirement already satisfied: decorator in /scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages (from tfp-nightly) (5.0.9)
Installing collected packages: tfp-nightly
  Attempting uninstall: tfp-nightly
    Found existing installation: tfp-nightly 0.14.0.dev20210525
    Uninstalling tfp-nightly-0.14.0.dev20210525:
      Successfully uninstalled tfp-nightly-0.14.0.dev20210525
Successfully installed tfp-nightly-0.14.0.dev20210527
(/scratch/gpfs/altosaar/environment-jax) 
# altosaar @ host in ~/projects/variational-autoencoder on git:inverse-autoregressive-flow x [16:40:02] C:1
$ python
Python 3.9.5 (default, May 18 2021, 19:34:48) 
[GCC 7.3.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from tensorflow_probability.substrates import jax as tfp
>>> tfb = tfp.bijectors
2021-05-27 16:40:22.317937: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/write/count
2021-05-27 16:40:22.318174: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/read/count
2021-05-27 16:40:22.318194: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/write/api
2021-05-27 16:40:22.318204: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/read/api
2021-05-27 16:40:23.948727: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/write/count
2021-05-27 16:40:23.948790: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/read/count
2021-05-27 16:40:23.948809: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/write/api
2021-05-27 16:40:23.948820: E tensorflow/core/lib/monitoring/collection_registry.cc:77] Cannot register 2 metrics with the same name: /tensorflow/core/saved_model/read/api
>>> tfb.AutoregressiveNetwork(params=2, hidden_units=[2])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/scratch/gpfs/altosaar/environment-jax/lib/python3.9/site-packages/tensorflow_probability/substrates/jax/bijectors/masked_autoregressive.py", line 981, in __init__
    self._kernel_constraint = tf.keras.constraints.get(kernel_constraint)
AttributeError: module 'tensorflow_probability.python.internal.backend.jax.keras' has no attribute 'constraints'
>>> 

Start at tensorflow_probability/substrates/jax/bijectors/masked_autoregressive.py around line 981, where tfb.AutoregressiveNetwork(params=2, hidden_units=[2]) fails while accessing keras.constraints. Reproduce the traceback in a Python environment with tf-nightly and tfp-nightly, then verify that constructing the network completes without the reported AttributeError.