code: def test_donation_vector(self):
      params = {"a": jnp.ones([]), "b": jnp.ones([])}
      state = {"c": jnp.ones([]), "d": jnp.ones([])}
      x = jnp.ones([])
      args = params, state, x
  
      for size in range(4):
          for donate_argnums in it.permutations((0, 1, 2), size):
              for kwargs in ({}, {"a": x}):
                  expected = ()
                  expected += (
                      (True, True)
                      if 0 in donate_argnums
                      else (False, False)
                  )
                  expected += (
                      (True, True)
                      if 1 in donate_argnums
                      else (False, False)
                  )
                  expected += (
                      (True,) if 2 in donate_argnums else (False,)
                  )
                  if kwargs:
                      expected += (False,)
                  self.assertEqual(
                      expected,
                      api_util.donation_vector(
                          donate_argnums, (), args, kwargs
                      ),
                  )
  
file_path: jax/tests/api_util_test.py
imports: import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  from jax._src import api_util
  from jax import numpy
  from jax._src import test_util
  from jax import config
label_desc: Use image node_id 1 for calling the ApiUtilTest obj's underlying member method code with example usage: obj.test_donation_vector() without return types