class_docstr: Masked Autoencoder for Distribution Estimation [Germain et al. (2015)][1].
  
  A `AutoregressiveNetwork` takes as input a Tensor of shape `[..., event_size]`
  and returns a Tensor of shape `[..., event_size, params]`.
  
  The output satisfies the autoregressive property.  That is, the layer is
  configured with some permutation `ord` of `{0, ..., event_size-1}` (i.e., an
  ordering of the input dimensions), and the output `output[batch_idx, i, ...]`
  for input dimension `i` depends only on inputs `x[batch_idx, j]` where
  `ord(j) < ord(i)`.  The autoregressive property allows us to use
  `output[batch_idx, i]` to parameterize conditional distributions:
    `p(x[batch_idx, i] | x[batch_idx, j] for ord(j) < ord(i))`
  which give us a tractable distribution over input `x[batch_idx]`:
    `p(x[batch_idx]) = prod_i p(x[batch_idx, ord(i)] | x[batch_idx, ord(0:i)])`
  
  For example, when `params` is 2, the output of the layer can parameterize
  the location and log-scale of an autoregressive Gaussian distribution.
  
  #### Example
  
  The `AutoregressiveNetwork` can be used to do density estimation as is shown
  in the below example:
  
  ```python
  # Generate data -- as in Figure 1 in [Papamakarios et al. (2017)][2]).
  n = 2000
  x2 = np.random.randn(n).astype(dtype=np.float32) * 2.
  x1 = np.random.randn(n).astype(dtype=np.float32) + (x2 * x2 / 4.)
  data = np.stack([x1, x2], axis=-1)
  
  # Density estimation with MADE.
  made = tfb.AutoregressiveNetwork(params=2, hidden_units=[10, 10])
  
  distribution = tfd.TransformedDistribution(
      distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[2]),
      bijector=tfb.MaskedAutoregressiveFlow(made))
  
  # Construct and fit model.
  x_ = tfkl.Input(shape=(2,), dtype=tf.float32)
  log_prob_ = distribution.log_prob(x_)
  model = tfk.Model(x_, log_prob_)
  
  model.compile(optimizer=tf_keras.optimizers.Adam(),
                loss=lambda _, log_prob: -log_prob)
  
  batch_size = 25
  model.fit(x=data,
            y=np.zeros((n, 0), dtype=np.float32),
            batch_size=batch_size,
            epochs=1,
            steps_per_epoch=1,  # Usually `n // batch_size`.
            shuffle=True,
            verbose=True)
  
  # Use the fitted distribution.
  distribution.sample((3, 1))
  distribution.log_prob(np.ones((3, 2), dtype=np.float32))
  ```
  
  The `conditional` argument can be used to instead build a conditional density
  estimator. To do this the conditioning variable must be passed as a `kwarg`:
  
  ```python
  # Generate data as the mixture of two distributions.
  n = 2000
  c = np.r_[
    np.zeros(n//2),
    np.ones(n//2)
  ]
  mean_0, mean_1 = 0, 5
  x = np.r_[
    np.random.randn(n//2).astype(dtype=np.float32) + mean_0,
    np.random.randn(n//2).astype(dtype=np.float32) + mean_1
  ]
  
  # Density estimation with MADE.
  made = tfb.AutoregressiveNetwork(
    params=2,
    hidden_units=[2, 2],
    event_shape=(1,),
    conditional=True,
    kernel_initializer=tfk.initializers.VarianceScaling(0.1),
    conditional_event_shape=(1,)
  )
  
  distribution = tfd.TransformedDistribution(
    distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[1]),
    bijector=tfb.MaskedAutoregressiveFlow(made))
  
  # Construct and fit model.
  x_ = tfkl.Input(shape=(1,), dtype=tf.float32)
  c_ = tfkl.Input(shape=(1,), dtype=tf.float32)
  log_prob_ = distribution.log_prob(
    x_, bijector_kwargs={'conditional_input': c_})
  model = tfk.Model([x_, c_], log_prob_)
  
  model.compile(optimizer=tf_keras.optimizers.Adam(learning_rate=0.1),
                loss=lambda _, log_prob: -log_prob)
  
  batch_size = 25
  model.fit(x=[x, c],
            y=np.zeros((n, 0), dtype=np.float32),
            batch_size=batch_size,
            epochs=3,
            steps_per_epoch=n // batch_size,
            shuffle=True,
            verbose=True)
  
  # Use the fitted distribution to sample condition on c = 1
  n_samples = 1000
  cond = 1
  samples = distribution.sample(
    (n_samples,),
    bijector_kwargs={'conditional_input': cond * np.ones((n_samples, 1))})
  ```
  
  #### Examples: Handling Rank-2+ Tensors
  
  `AutoregressiveNetwork` can be used as a building block to achieve different
  autoregressive structures over rank-2+ tensors.  For example, suppose we want
  to build an autoregressive distribution over images with dimension `[weight,
  height, channels]` with `channels = 3`:
  
   1. We can parameterize a 'fully autoregressive' distribution, with
      cross-channel and within-pixel autoregressivity:
      ```
          r0    g0   b0     r0    g0   b0       r0   g0    b0
          ^   ^      ^         ^   ^   ^         ^      ^   ^
          |  /  ____/           \  |  /           \____  \  |
          | /__/                 \ | /                 \__\ |
          r1    g1   b1     r1 <- g1   b1       r1   g1 <- b1
                                               ^          |
                                                \_________/
      ```
  
      as:
      ```python
      # Generate random images for training data.
      images = np.random.uniform(size=(100, 8, 8, 3)).astype(np.float32)
      n, width, height, channels = images.shape
  
      # Reshape images to achieve desired autoregressivity.
      event_shape = [height * width * channels]
      reshaped_images = tf.reshape(images, [n, event_shape])
  
      # Density estimation with MADE.
      made = tfb.AutoregressiveNetwork(params=2, event_shape=event_shape,
                                       hidden_units=[20, 20], activation='relu')
      distribution = tfd.TransformedDistribution(
      distribution=tfd.Sample(
          tfd.Normal(loc=0., scale=1.), sample_shape=[dims]),
      bijector=tfb.MaskedAutoregressiveFlow(made))
  
      # Construct and fit model.
      x_ = tfkl.Input(shape=event_shape, dtype=tf.float32)
      log_prob_ = distribution.log_prob(x_)
      model = tfk.Model(x_, log_prob_)
  
      model.compile(optimizer=tf_keras.optimizers.Adam(),
                    loss=lambda _, log_prob: -log_prob)
  
      batch_size = 10
      model.fit(x=data,
                y=np.zeros((n, 0), dtype=np.float32),
                batch_size=batch_size,
                epochs=10,
                steps_per_epoch=n // batch_size,
                shuffle=True,
                verbose=True)
  
      # Use the fitted distribution.
      distribution.sample((3, 1))
      distribution.log_prob(np.ones((5, 8, 8, 3), dtype=np.float32))
      ```
  
   2. We can parameterize a distribution with neither cross-channel nor
      within-pixel autoregressivity:
      ```
          r0    g0   b0
          ^     ^    ^
          |     |    |
          |     |    |
          r1    g1   b1
      ```
  
      as:
      ```python
      # Generate fake images.
      images = np.random.choice([0, 1], size=(100, 8, 8, 3))
      n, width, height, channels = images.shape
  
      # Reshape images to achieve desired autoregressivity.
      reshaped_images = np.transpose(
          np.reshape(images, [n, width * height, channels]),
          axes=[0, 2, 1])
  
      made = tfb.AutoregressiveNetwork(params=1, event_shape=[width * height],
                                       hidden_units=[20, 20], activation='relu')
  
      # Density estimation with MADE.
      #
      # NOTE: Parameterize an autoregressive distribution over an event_shape of
      # [channels, width * height], with univariate Bernoulli conditional
      # distributions.
      distribution = tfd.Autoregressive(
          lambda x: tfd.Independent(
              tfd.Bernoulli(logits=tf.unstack(made(x), axis=-1)[0],
                            dtype=tf.float32),
              reinterpreted_batch_ndims=2),
          sample0=tf.zeros([channels, width * height], dtype=tf.float32))
  
      # Construct and fit model.
      x_ = tfkl.Input(shape=(channels, width * height), dtype=tf.float32)
      log_prob_ = distribution.log_prob(x_)
      model = tfk.Model(x_, log_prob_)
  
      model.compile(optimizer=tf_keras.optimizers.Adam(),
                    loss=lambda _, log_prob: -log_prob)
  
      batch_size = 10
      model.fit(x=reshaped_images,
                y=np.zeros((n, 0), dtype=np.float32),
                batch_size=batch_size,
                epochs=10,
                steps_per_epoch=n // batch_size,
                shuffle=True,
                verbose=True)
  
      distribution.sample(7)
      distribution.log_prob(np.ones((4, 8, 8, 3), dtype=np.float32))
      ```
  
      Note that one set of weights is shared for the mapping for each channel
      from image to distribution parameters -- i.e., the mapping
      `layer(reshaped_images[..., channel, :])`, where `channel` is 0, 1, or 2.
  
      To use separate weights for each channel, we could construct an
      `AutoregressiveNetwork` and `TransformedDistribution` for each channel,
      and combine them with a `tfd.Blockwise` distribution.
  
  #### References
  
  [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
       Masked Autoencoder for Distribution Estimation. In _International
       Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509
  
  [2]: George Papamakarios, Theo Pavlakou, Iain Murray, Masked Autoregressive
       Flow for Density Estimation.  In _Neural Information Processing Systems_,
       2017. https://arxiv.org/abs/1705.07057
class_docstr_tok: ["Masked","Autoencoder","for","Distribution","Estimation","[","Germain","et","al",".","(","2015",")","]","[","1","]",".","A","`","AutoregressiveNetwork","`","takes","as","input","a","Tensor","of","shape","`","[","...",",","event_size","]","`","and","returns","a","Tensor","of","shape","`","[","...",",","event_size",",","params","]","`",".","The","output","satisfies","the","autoregressive","property",".","That","is",",","the","layer","is","configured","with","some","permutation","`","ord","`","of","`","{","0",",","...",",","event_size-1","}","`","(","i.e.",",","an","ordering","of","the","input","dimensions",")",",","and","the","output","`","output","[","batch_idx",",","i",",","...","]","`","for","input","dimension","`","i","`","depends","only","on","inputs","`","x","[","batch_idx",",","j","]","`","where","`","ord","(","j",")","<","ord","(","i",")","`",".","The","autoregressive","property","allows","us","to","use","`","output","[","batch_idx",",","i","]","`","to","parameterize","conditional","distributions",":","`","p","(","x","[","batch_idx",",","i","]","|","x","[","batch_idx",",","j","]","for","ord","(","j",")","<","ord","(","i",")",")","`","which","give","us","a","tractable","distribution","over","input","`","x","[","batch_idx","]","`",":","`","p","(","x","[","batch_idx","]",")","=","prod_i","p","(","x","[","batch_idx",",","ord","(","i",")","]","|","x","[","batch_idx",",","ord","(","0",":","i",")","]",")","`","For","example",",","when","`","params","`","is","2",",","the","output","of","the","layer","can","parameterize","the","location","and","log-scale","of","an","autoregressive","Gaussian","distribution",".","#","#","#","#","Example","The","`","AutoregressiveNetwork","`","can","be","used","to","do","density","estimation","as","is","shown","in","the","below","example",":","``","`","python","#","Generate","data","--","as","in","Figure","1","in","[","Papamakarios","et","al",".","(","2017",")","]","[","2","]",")",".","n","=","2000","x2","=","np.random.randn","(","n",")",".astype","(","dtype=np.float32",")","*","2.","x1","=","np.random.randn","(","n",")",".astype","(","dtype=np.float32",")","+","(","x2","*","x2","\/","4",".",")","data","=","np.stack","(","[","x1",",","x2","]",",","axis=-1",")","#","Density","estimation","with","MADE",".","made","=","tfb.AutoregressiveNetwork","(","params=2",",","hidden_units=","[","10",",","10","]",")","distribution","=","tfd.TransformedDistribution","(","distribution=tfd.Sample","(","tfd.Normal","(","loc=0.",",","scale=1",".",")",",","sample_shape=","[","2","]",")",",","bijector=tfb.MaskedAutoregressiveFlow","(","made",")",")","#","Construct","and","fit","model",".","x_","=","tfkl.Input","(","shape=","(","2",",",")",",","dtype=tf.float32",")","log_prob_","=","distribution.log_prob","(","x_",")","model","=","tfk.Model","(","x_",",","log_prob_",")","model.compile","(","optimizer=tf_keras.optimizers.Adam","(",")",",","loss=lambda","_",",","log_prob",":","-log_prob",")","batch_size","=","25","model.fit","(","x=data",",","y=np.zeros","(","(","n",",","0",")",",","dtype=np.float32",")",",","batch_size=batch_size",",","epochs=1",",","steps_per_epoch=1",",","#","Usually","`","n","\/\/","batch_size","`",".","shuffle=True",",","verbose=True",")","#","Use","the","fitted","distribution",".","distribution.sample","(","(","3",",","1",")",")","distribution.log_prob","(","np.ones","(","(","3",",","2",")",",","dtype=np.float32",")",")","``","`","The","`","conditional","`","argument","can","be","used","to","instead","build","a","conditional","density","estimator",".","To","do","this","the","conditioning","variable","must","be","passed","as","a","`","kwarg","`",":","``","`","python","#","Generate","data","as","the","mixture","of","two","distributions",".","n","=","2000","c","=","np.r_","[","np.zeros","(","n\/\/2",")",",","np.ones","(","n\/\/2",")","]","mean_0",",","mean_1","=","0",",","5","x","=","np.r_","[","np.random.randn","(","n\/\/2",")",".astype","(","dtype=np.float32",")","+","mean_0",",","np.random.randn","(","n\/\/2",")",".astype","(","dtype=np.float32",")","+","mean_1","]","#","Density","estimation","with","MADE",".","made","=","tfb.AutoregressiveNetwork","(","params=2",",","hidden_units=","[","2",",","2","]",",","event_shape=","(","1",",",")",",","conditional=True",",","kernel_initializer=tfk.initializers.VarianceScaling","(","0.1",")",",","conditional_event_shape=","(","1",",",")",")","distribution","=","tfd.TransformedDistribution","(","distribution=tfd.Sample","(","tfd.Normal","(","loc=0.",",","scale=1",".",")",",","sample_shape=","[","1","]",")",",","bijector=tfb.MaskedAutoregressiveFlow","(","made",")",")","#","Construct","and","fit","model",".","x_","=","tfkl.Input","(","shape=","(","1",",",")",",","dtype=tf.float32",")","c_","=","tfkl.Input","(","shape=","(","1",",",")",",","dtype=tf.float32",")","log_prob_","=","distribution.log_prob","(","x_",",","bijector_kwargs=","{","'conditional_input","'",":","c_","}",")","model","=","tfk.Model","(","[","x_",",","c_","]",",","log_prob_",")","model.compile","(","optimizer=tf_keras.optimizers.Adam","(","learning_rate=0.1",")",",","loss=lambda","_",",","log_prob",":","-log_prob",")","batch_size","=","25","model.fit","(","x=","[","x",",","c","]",",","y=np.zeros","(","(","n",",","0",")",",","dtype=np.float32",")",",","batch_size=batch_size",",","epochs=3",",","steps_per_epoch=n","\/\/","batch_size",",","shuffle=True",",","verbose=True",")","#","Use","the","fitted","distribution","to","sample","condition","on","c","=","1","n_samples","=","1000","cond","=","1","samples","=","distribution.sample","(","(","n_samples",",",")",",","bijector_kwargs=","{","'conditional_input","'",":","cond","*","np.ones","(","(","n_samples",",","1",")",")","}",")","``","`","#","#","#","#","Examples",":","Handling","Rank-2+","Tensors","`","AutoregressiveNetwork","`","can","be","used","as","a","building","block","to","achieve","different","autoregressive","structures","over","rank-2+","tensors",".","For","example",",","suppose","we","want","to","build","an","autoregressive","distribution","over","images","with","dimension","`","[","weight",",","height",",","channels","]","`","with","`","channels","=","3","`",":","1",".","We","can","parameterize","a","'fully","autoregressive","'","distribution",",","with","cross-channel","and","within-pixel","autoregressivity",":","``","`","r0","g0","b0","r0","g0","b0","r0","g0","b0","^","^","^","^","^","^","^","^","^","|","\/","____\/","\\","|","\/","\\____","\\","|","|","\/__\/","\\","|","\/","\\__\\","|","r1","g1","b1","r1","<","-","g1","b1","r1","g1","<","-","b1","^","|","\\_________\/","``","`","as",":","``","`","python","#","Generate","random","images","for","training","data",".","images","=","np.random.uniform","(","size=","(","100",",","8",",","8",",","3",")",")",".astype","(","np.float32",")","n",",","width",",","height",",","channels","=","images.shape","#","Reshape","images","to","achieve","desired","autoregressivity",".","event_shape","=","[","height","*","width","*","channels","]","reshaped_images","=","tf.reshape","(","images",",","[","n",",","event_shape","]",")","#","Density","estimation","with","MADE",".","made","=","tfb.AutoregressiveNetwork","(","params=2",",","event_shape=event_shape",",","hidden_units=","[","20",",","20","]",",","activation='relu","'",")","distribution","=","tfd.TransformedDistribution","(","distribution=tfd.Sample","(","tfd.Normal","(","loc=0.",",","scale=1",".",")",",","sample_shape=","[","dims","]",")",",","bijector=tfb.MaskedAutoregressiveFlow","(","made",")",")","#","Construct","and","fit","model",".","x_","=","tfkl.Input","(","shape=event_shape",",","dtype=tf.float32",")","log_prob_","=","distribution.log_prob","(","x_",")","model","=","tfk.Model","(","x_",",","log_prob_",")","model.compile","(","optimizer=tf_keras.optimizers.Adam","(",")",",","loss=lambda","_",",","log_prob",":","-log_prob",")","batch_size","=","10","model.fit","(","x=data",",","y=np.zeros","(","(","n",",","0",")",",","dtype=np.float32",")",",","batch_size=batch_size",",","epochs=10",",","steps_per_epoch=n","\/\/","batch_size",",","shuffle=True",",","verbose=True",")","#","Use","the","fitted","distribution",".","distribution.sample","(","(","3",",","1",")",")","distribution.log_prob","(","np.ones","(","(","5",",","8",",","8",",","3",")",",","dtype=np.float32",")",")","``","`","2",".","We","can","parameterize","a","distribution","with","neither","cross-channel","nor","within-pixel","autoregressivity",":","``","`","r0","g0","b0","^","^","^","|","|","|","|","|","|","r1","g1","b1","``","`","as",":","``","`","python","#","Generate","fake","images",".","images","=","np.random.choice","(","[","0",",","1","]",",","size=","(","100",",","8",",","8",",","3",")",")","n",",","width",",","height",",","channels","=","images.shape","#","Reshape","images","to","achieve","desired","autoregressivity",".","reshaped_images","=","np.transpose","(","np.reshape","(","images",",","[","n",",","width","*","height",",","channels","]",")",",","axes=","[","0",",","2",",","1","]",")","made","=","tfb.AutoregressiveNetwork","(","params=1",",","event_shape=","[","width","*","height","]",",","hidden_units=","[","20",",","20","]",",","activation='relu","'",")","#","Density","estimation","with","MADE",".","#","#","NOTE",":","Parameterize","an","autoregressive","distribution","over","an","event_shape","of","#","[","channels",",","width","*","height","]",",","with","univariate","Bernoulli","conditional","#","distributions",".","distribution","=","tfd.Autoregressive","(","lambda","x",":","tfd.Independent","(","tfd.Bernoulli","(","logits=tf.unstack","(","made","(","x",")",",","axis=-1",")","[","0","]",",","dtype=tf.float32",")",",","reinterpreted_batch_ndims=2",")",",","sample0=tf.zeros","(","[","channels",",","width","*","height","]",",","dtype=tf.float32",")",")","#","Construct","and","fit","model",".","x_","=","tfkl.Input","(","shape=","(","channels",",","width","*","height",")",",","dtype=tf.float32",")","log_prob_","=","distribution.log_prob","(","x_",")","model","=","tfk.Model","(","x_",",","log_prob_",")","model.compile","(","optimizer=tf_keras.optimizers.Adam","(",")",",","loss=lambda","_",",","log_prob",":","-log_prob",")","batch_size","=","10","model.fit","(","x=reshaped_images",",","y=np.zeros","(","(","n",",","0",")",",","dtype=np.float32",")",",","batch_size=batch_size",",","epochs=10",",","steps_per_epoch=n","\/\/","batch_size",",","shuffle=True",",","verbose=True",")","distribution.sample","(","7",")","distribution.log_prob","(","np.ones","(","(","4",",","8",",","8",",","3",")",",","dtype=np.float32",")",")","``","`","Note","that","one","set","of","weights","is","shared","for","the","mapping","for","each","channel","from","image","to","distribution","parameters","--","i.e.",",","the","mapping","`","layer","(","reshaped_images","[","...",",","channel",",",":","]",")","`",",","where","`","channel","`","is","0",",","1",",","or","2",".","To","use","separate","weights","for","each","channel",",","we","could","construct","an","`","AutoregressiveNetwork","`","and","`","TransformedDistribution","`","for","each","channel",",","and","combine","them","with","a","`","tfd.Blockwise","`","distribution",".","#","#","#","#","References","[","1","]",":","Mathieu","Germain",",","Karol","Gregor",",","Iain","Murray",",","and","Hugo","Larochelle",".","MADE",":","Masked","Autoencoder","for","Distribution","Estimation",".","In","_International","Conference","on","Machine","Learning_",",","2015.","https",":","\/\/arxiv.org\/abs\/1502.03509","[","2","]",":","George","Papamakarios",",","Theo","Pavlakou",",","Iain","Murray",",","Masked","Autoregressive","Flow","for","Density","Estimation",".","In","_Neural","Information","Processing","Systems_",",","2017.","https",":","\/\/arxiv.org\/abs\/1705.07057"]
docstr: Transforms the inputs and returns the outputs.
  
  Suppose `x` has shape `batch_shape + event_shape` and `conditional_input`
  has shape `conditional_batch_shape + conditional_event_shape`. Then, the
  output shape is:
  `broadcast(batch_shape, conditional_batch_shape) + event_shape + [params]`.
  
  Also see `tfkl.Layer.call` for some generic discussion about Layer calling.
  
  Args:
    x: A `Tensor`. Primary input to the layer.
    conditional_input: A `Tensor. Conditional input to the layer. This is
      required iff the layer is conditional.
  
  Returns:
    y: A `Tensor`. The output of the layer. Note that the leading dimensions
       follow broadcasting rules described above.
docstr_tok: ["Transforms","the","inputs","and","returns","the","outputs",".","Suppose","`","x","`","has","shape","`","batch_shape","+","event_shape","`","and","`","conditional_input","`","has","shape","`","conditional_batch_shape","+","conditional_event_shape","`",".","Then",",","the","output","shape","is",":","`","broadcast","(","batch_shape",",","conditional_batch_shape",")","+","event_shape","+","[","params","]","`",".","Also","see","`","tfkl.Layer.call","`","for","some","generic","discussion","about","Layer","calling",".","Args",":","x",":","A","`","Tensor","`",".","Primary","input","to","the","layer",".","conditional_input",":","A","`","Tensor",".","Conditional","input","to","the","layer",".","This","is","required","iff","the","layer","is","conditional",".","Returns",":","y",":","A","`","Tensor","`",".","The","output","of","the","layer",".","Note","that","the","leading","dimensions","follow","broadcasting","rules","described","above","."]
code: def call(self, x, conditional_input=None):
      """Transforms the inputs and returns the outputs.
  
      Suppose `x` has shape `batch_shape + event_shape` and `conditional_input`
      has shape `conditional_batch_shape + conditional_event_shape`. Then, the
      output shape is:
      `broadcast(batch_shape, conditional_batch_shape) + event_shape + [params]`.
  
      Also see `tfkl.Layer.call` for some generic discussion about Layer calling.
  
      Args:
        x: A `Tensor`. Primary input to the layer.
        conditional_input: A `Tensor. Conditional input to the layer. This is
          required iff the layer is conditional.
  
      Returns:
        y: A `Tensor`. The output of the layer. Note that the leading dimensions
           follow broadcasting rules described above.
      """
      with tf.name_scope(self.name or "AutoregressiveNetwork_call"):
          x = tf.convert_to_tensor(x, dtype=self.dtype, name="x")
          # TODO(b/67594795): Better support for dynamic shapes.
          input_shape = ps.shape(x)
          if tensorshape_util.rank(x.shape) == 1:
              x = x[tf.newaxis, ...]
          if self._conditional:
              if conditional_input is None:
                  raise ValueError(
                      "`conditional_input` must be passed as a named "
                      "argument"
                  )
              conditional_input = tf.convert_to_tensor(
                  conditional_input,
                  dtype=self.dtype,
                  name="conditional_input",
              )
              conditional_batch_shape = ps.shape(conditional_input)[:-1]
              if tensorshape_util.rank(conditional_input.shape) == 1:
                  conditional_input = conditional_input[tf.newaxis, ...]
              x = [x, conditional_input]
              output_shape = ps.concat(
                  [
                      ps.broadcast_shape(
                          conditional_batch_shape, input_shape[:-1]
                      ),
                      input_shape[-1:],
                  ],
                  axis=0,
              )
          else:
              output_shape = input_shape
          return tf.reshape(
              self._network(x),
              tf.concat([output_shape, [self._params]], axis=0),
          )
  
imports: import numpy
  import six
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.math.numeric import clip_by_value_preserve_gradient
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 3 for calling the AutoregressiveNetwork obj's underlying member method code with example usage: obj.call(x, conditional_input) and returns: tf