code: def test_conditional_compose_transform_module(
      batch_shape, input_dim, context_dim, cache_size
  ):
      conditional_transforms = [
          T.AffineTransform(1.0, 2.0),
          T.Spline(input_dim),
          T.conditional_spline(input_dim, context_dim, [5]),
          T.SoftplusTransform(),
          T.conditional_spline(input_dim, context_dim, [6]),
      ]
      cond_transform = (
          dist.conditional.ConditionalComposeTransformModule(
              conditional_transforms, cache_size=cache_size
          )
      )
  
      base_dist = (
          dist.Normal(0, 1)
          .expand(batch_shape + (input_dim,))
          .to_event(1)
      )
      cond_dist = dist.ConditionalTransformedDistribution(
          base_dist, [cond_transform]
      )
  
      context = torch.rand(batch_shape + (context_dim,))
      d = cond_dist.condition(context)
      transform = d.transforms[0]
      assert isinstance(transform, T.ComposeTransformModule)
  
      data = d.rsample()
      assert data.shape == batch_shape + (input_dim,)
      assert d.log_prob(data).shape == batch_shape
  
      actual_params = set(cond_transform.parameters())
      expected_params = set(
          torch.nn.ModuleList(
              [
                  t
                  for t in conditional_transforms
                  if isinstance(t, torch.nn.Module)
              ]
          ).parameters()
      )
      assert set() != actual_params == expected_params
  
      noise = base_dist.rsample()
      expected = noise
      for t in conditional_transforms:
          expected = (
              t.condition(context) if hasattr(t, "condition") else t
          )(expected)
  
      actual = transform(noise)
      assert_close(actual, expected)
  
      assert_close(cond_transform.inv.condition(context)(actual), noise)
      assert_close(
          cond_transform.condition(context).inv(expected), noise
      )
  
imports: import operator
  from functools import partial, reduce
  from unittest import TestCase
  import pytest
  import torch
  import pyro.distributions
  import pyro.distributions.transforms
  from pyro.distributions import constraints
  from tests.common import assert_close
label_desc: Use image node_id 4 for calling a global function with example usage: test_conditional_compose_transform_module(batch_shape, input_dim, context_dim, cache_size) without return types