text: """Implementations of autoregressive transforms."""
  
  import numpy as np
  import torch
  from torch.nn import functional as F
  
  from neural_spline_flows import utils
  from neural_spline_flows.nde import transforms
  from neural_spline_flows.nde.transforms import made as made_module
  from neural_spline_flows.nde.transforms import splines
  
  class AutoregressiveTransform(transforms.Transform):
      """Transforms each input variable with an invertible elementwise transformation.
  
      The parameters of each invertible elementwise transformation can be functions of previous input
      variables, but they must not depend on the current or any following input variables.
  
      NOTE: Calculating the inverse transform is D times slower than calculating the
      forward transform, where D is the dimensionality of the input to the transform.
      """
      def __init__(self, autoregressive_net):
          super(AutoregressiveTransform, self).__init__()
          self.autoregressive_net = autoregressive_net
  
      def forward(self, inputs, context=None):
          autoregressive_params = self.autoregressive_net(inputs, context)
          outputs, logabsdet = self._elementwise_forward(inputs, autoregressive_params)
          return outputs, logabsdet
  
      def inverse(self, inputs, context=None):
          num_inputs = np.prod(inputs.shape[1:])
          outputs = torch.zeros_like(inputs)
          logabsdet = None
          for _ in range(num_inputs):
              autoregressive_params = self.autoregressive_net(outputs, context)
              outputs, logabsdet = self._elementwise_inverse(inputs, autoregressive_params)
          return outputs, logabsdet
  
      def _output_dim_multiplier(self):
          raise NotImplementedError()
  
      def _elementwise_forward(self, inputs, autoregressive_params):
          raise NotImplementedError()
  
      def _elementwise_inverse(self, inputs, autoregressive_params):
          raise NotImplementedError()
  
  
  class MaskedAffineAutoregressiveTransform(AutoregressiveTransform):
      def __init__(self,
                   features,
                   hidden_features,
                   context_features=None,
                   num_blocks=2,
                   use_residual_blocks=True,
                   random_mask=False,
                   activation=F.relu,
                   dropout_probability=0.,
                   use_batch_norm=False):
          self.features = features
          made = made_module.MADE(
              features=features,
              hidden_features=hidden_features,
              context_features=context_features,
              num_blocks=num_blocks,
              output_multiplier=self._output_dim_multiplier(),
              use_residual_blocks=use_residual_blocks,
              random_mask=random_mask,
              activation=activation,
              dropout_probability=dropout_probability,
              use_batch_norm=use_batch_norm,
          )
          super(MaskedAffineAutoregressiveTransform, self).__init__(made)
  
      def _output_dim_multiplier(self):
          return 2
  
      def _elementwise_forward(self, inputs, autoregressive_params):
          unconstrained_scale, shift = self._unconstrained_scale_and_shift(autoregressive_params)
          scale = torch.sigmoid(unconstrained_scale + 2.) + 1e-3
          log_scale = torch.log(scale)
          outputs = scale * inputs + shift
          logabsdet = utils.sum_except_batch(log_scale, num_batch_dims=1)
          return outputs, logabsdet
  
      def _elementwise_inverse(self, inputs, autoregressive_params):
          unconstrained_scale, shift = self._unconstrained_scale_and_shift(autoregressive_params)
          scale = torch.sigmoid(unconstrained_scale + 2.) + 1e-3
          log_scale = torch.log(scale)
          outputs = (inputs - shift) / scale
          logabsdet = -utils.sum_except_batch(log_scale, num_batch_dims=1)
          return outputs, logabsdet
  
      def _unconstrained_scale_and_shift(self, autoregressive_params):
          # split_idx = autoregressive_params.size(1) // 2
          # unconstrained_scale = autoregressive_params[..., :split_idx]
          # shift = autoregressive_params[..., split_idx:]
          # return unconstrained_scale, shift
          autoregressive_params = autoregressive_params.view(
              -1, self.features, self._output_dim_multiplier()
          )
          unconstrained_scale = autoregressive_params[..., 0]
          shift = autoregressive_params[..., 1]
          return unconstrained_scale, shift
  
  
  class MaskedPiecewiseLinearAutoregressiveTransform(AutoregressiveTransform):
      def __init__(self,
                   num_bins,
                   features,
                   hidden_features,
                   context_features=None,
                   num_blocks=2,
                   use_residual_blocks=True,
                   random_mask=False,
                   activation=F.relu,
                   dropout_probability=0.,
                   use_batch_norm=False):
          self.num_bins = num_bins
          self.features = features
          made = made_module.MADE(
              features=features,
              hidden_features=hidden_features,
              context_features=context_features,
              num_blocks=num_blocks,
              output_multiplier=self._output_dim_multiplier(),
              use_residual_blocks=use_residual_blocks,
              random_mask=random_mask,
              activation=activation,
              dropout_probability=dropout_probability,
              use_batch_norm=use_batch_norm,
          )
          super().__init__(made)
  
      def _output_dim_multiplier(self):
          return self.num_bins
  
      def _elementwise(self, inputs, autoregressive_params, inverse=False):
          batch_size = inputs.shape[0]
  
          unnormalized_pdf = autoregressive_params.view(batch_size,
                                                        self.features,
                                                        self._output_dim_multiplier())
  
          outputs, logabsdet = splines.linear_spline(inputs=inputs,
                                                     unnormalized_pdf=unnormalized_pdf,
                                                     inverse=inverse)
  
          return outputs, utils.sum_except_batch(logabsdet)
  
      def _elementwise_forward(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params)
  
      def _elementwise_inverse(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params, inverse=True)
  
  
  class MaskedPiecewiseQuadraticAutoregressiveTransform(AutoregressiveTransform):
      def __init__(self,
                   features,
                   hidden_features,
                   context_features=None,
                   num_bins=10,
                   num_blocks=2,
                   tails=None,
                   tail_bound=1.,
                   use_residual_blocks=True,
                   random_mask=False,
                   activation=F.relu,
                   dropout_probability=0.,
                   use_batch_norm=False,
                   min_bin_width=splines.rational_quadratic.DEFAULT_MIN_BIN_WIDTH,
                   min_bin_height=splines.rational_quadratic.DEFAULT_MIN_BIN_HEIGHT,
                   min_derivative=splines.rational_quadratic.DEFAULT_MIN_DERIVATIVE
                   ):
          self.num_bins = num_bins
          self.min_bin_width = min_bin_width
          self.min_bin_height = min_bin_height
          self.min_derivative = min_derivative
          self.tails = tails
          self.tail_bound = tail_bound
          self.features = features
          made = made_module.MADE(
              features=features,
              hidden_features=hidden_features,
              context_features=context_features,
              num_blocks=num_blocks,
              output_multiplier=self._output_dim_multiplier(),
              use_residual_blocks=use_residual_blocks,
              random_mask=random_mask,
              activation=activation,
              dropout_probability=dropout_probability,
              use_batch_norm=use_batch_norm,
          )
          super().__init__(made)
  
      def _output_dim_multiplier(self):
          if self.tails == 'linear':
              return self.num_bins * 2 - 1
          else:
              return self.num_bins * 2 + 1
  
      def _elementwise(self, inputs, autoregressive_params, inverse=False):
          batch_size = inputs.shape[0]
  
          transform_params = autoregressive_params.view(batch_size,
                                                        self.features,
                                                        self._output_dim_multiplier())
  
          unnormalized_widths = transform_params[..., :self.num_bins]
          unnormalized_heights = transform_params[..., self.num_bins:]
  
          if hasattr(self.autoregressive_net, 'hidden_features'):
              unnormalized_widths /= np.sqrt(self.autoregressive_net.hidden_features)
              unnormalized_heights /= np.sqrt(self.autoregressive_net.hidden_features)
  
          if self.tails is None:
              spline_fn = splines.quadratic_spline
              spline_kwargs = {}
          elif self.tails == 'linear':
              spline_fn = splines.unconstrained_quadratic_spline
              spline_kwargs = {
                  'tails': self.tails,
                  'tail_bound': self.tail_bound
              }
          else:
              raise ValueError
  
          outputs, logabsdet = spline_fn(
              inputs=inputs,
              unnormalized_heights=unnormalized_heights,
              unnormalized_widths=unnormalized_widths,
              inverse=inverse,
              min_bin_width=self.min_bin_width,
              min_bin_height=self.min_bin_height,
              **spline_kwargs
          )
  
          return outputs, utils.sum_except_batch(logabsdet)
  
      def _elementwise_forward(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params)
  
      def _elementwise_inverse(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params, inverse=True)
  
  
  class MaskedPiecewiseCubicAutoregressiveTransform(AutoregressiveTransform):
      def __init__(self,
                   num_bins,
                   features,
                   hidden_features,
                   context_features=None,
                   num_blocks=2,
                   use_residual_blocks=True,
                   random_mask=False,
                   activation=F.relu,
                   dropout_probability=0.,
                   use_batch_norm=False):
          self.num_bins = num_bins
          self.features = features
          made = made_module.MADE(
              features=features,
              hidden_features=hidden_features,
              context_features=context_features,
              num_blocks=num_blocks,
              output_multiplier=self._output_dim_multiplier(),
              use_residual_blocks=use_residual_blocks,
              random_mask=random_mask,
              activation=activation,
              dropout_probability=dropout_probability,
              use_batch_norm=use_batch_norm,
          )
          super(MaskedPiecewiseCubicAutoregressiveTransform, self).__init__(made)
  
      def _output_dim_multiplier(self):
          return self.num_bins * 2 + 2
  
      def _elementwise(self, inputs, autoregressive_params, inverse=False):
          batch_size = inputs.shape[0]
  
          transform_params = autoregressive_params.view(batch_size,
                                                   self.features,
                                                   self.num_bins * 2 + 2)
  
          unnormalized_widths = transform_params[...,:self.num_bins]
          unnormalized_heights = transform_params[...,self.num_bins:2*self.num_bins]
          derivatives = transform_params[...,2*self.num_bins:]
          unnorm_derivatives_left = derivatives[..., 0][..., None]
          unnorm_derivatives_right = derivatives[..., 1][..., None]
  
          if hasattr(self.autoregressive_net, 'hidden_features'):
              unnormalized_widths /= np.sqrt(self.autoregressive_net.hidden_features)
              unnormalized_heights /= np.sqrt(self.autoregressive_net.hidden_features)
  
          outputs, logabsdet = splines.cubic_spline(
              inputs=inputs,
              unnormalized_widths=unnormalized_widths,
              unnormalized_heights=unnormalized_heights,
              unnorm_derivatives_left=unnorm_derivatives_left,
              unnorm_derivatives_right=unnorm_derivatives_right,
              inverse=inverse
          )
          return outputs, utils.sum_except_batch(logabsdet)
  
      def _elementwise_forward(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params)
  
      def _elementwise_inverse(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params, inverse=True)
  
  
  class MaskedPiecewiseRationalQuadraticAutoregressiveTransform(AutoregressiveTransform):
      def __init__(self,
                   features,
                   hidden_features,
                   context_features=None,
                   num_bins=10,
                   tails=None,
                   tail_bound=1.,
                   num_blocks=2,
                   use_residual_blocks=True,
                   random_mask=False,
                   activation=F.relu,
                   dropout_probability=0.,
                   use_batch_norm=False,
                   min_bin_width=splines.rational_quadratic.DEFAULT_MIN_BIN_WIDTH,
                   min_bin_height=splines.rational_quadratic.DEFAULT_MIN_BIN_HEIGHT,
                   min_derivative=splines.rational_quadratic.DEFAULT_MIN_DERIVATIVE
                   ):
          self.num_bins = num_bins
          self.min_bin_width = min_bin_width
          self.min_bin_height = min_bin_height
          self.min_derivative = min_derivative
          self.tails = tails
          self.tail_bound = tail_bound
  
          autoregressive_net = made_module.MADE(
              features=features,
              hidden_features=hidden_features,
              context_features=context_features,
              num_blocks=num_blocks,
              output_multiplier=self._output_dim_multiplier(),
              use_residual_blocks=use_residual_blocks,
              random_mask=random_mask,
              activation=activation,
              dropout_probability=dropout_probability,
              use_batch_norm=use_batch_norm,
          )
  
          super().__init__(autoregressive_net)
  
      def _output_dim_multiplier(self):
          if self.tails == 'linear':
              return self.num_bins * 3 - 1
          elif self.tails is None:
              return self.num_bins * 3 + 1
          else:
              raise ValueError
  
      def _elementwise(self, inputs, autoregressive_params, inverse=False):
          batch_size, features = inputs.shape[0], inputs.shape[1]
  
          transform_params = autoregressive_params.view(
              batch_size,
              features,
              self._output_dim_multiplier()
          )
  
          unnormalized_widths = transform_params[...,:self.num_bins]
          unnormalized_heights = transform_params[...,self.num_bins:2*self.num_bins]
          unnormalized_derivatives = transform_params[...,2*self.num_bins:]
  
          if hasattr(self.autoregressive_net, 'hidden_features'):
              unnormalized_widths /= np.sqrt(self.autoregressive_net.hidden_features)
              unnormalized_heights /= np.sqrt(self.autoregressive_net.hidden_features)
  
          if self.tails is None:
              spline_fn = splines.rational_quadratic_spline
              spline_kwargs = {}
          elif self.tails == 'linear':
              spline_fn = splines.unconstrained_rational_quadratic_spline
              spline_kwargs = {
                  'tails': self.tails,
                  'tail_bound': self.tail_bound
              }
          else:
              raise ValueError
  
          outputs, logabsdet = spline_fn(
              inputs=inputs,
              unnormalized_widths=unnormalized_widths,
              unnormalized_heights=unnormalized_heights,
              unnormalized_derivatives=unnormalized_derivatives,
              inverse=inverse,
              min_bin_width=self.min_bin_width,
              min_bin_height=self.min_bin_height,
              min_derivative=self.min_derivative,
              **spline_kwargs
          )
  
          return outputs, utils.sum_except_batch(logabsdet)
  
      def _elementwise_forward(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params)
  
      def _elementwise_inverse(self, inputs, autoregressive_params):
          return self._elementwise(inputs, autoregressive_params, inverse=True)
  
  
  def main():
      inputs = torch.randn(16, 10)
      context = torch.randn(16, 24)
      transform = MaskedPiecewiseQuadraticAutoregressiveTransform(
          features=10,
          hidden_features=32,
          context_features=24,
          num_bins=10,
          tails='linear',
          num_blocks=2
      )
      outputs, logabsdet = transform(inputs, context)
      print(outputs.shape)
  
  
  if __name__ == '__main__':
      main()
  