code: def test_static_argnames(self, dtype, bs):
      # In this test, leave static_argnums at the default value too, in order to
      # test the default/most canonical path where `params` are the first arg.
      @partial(jax_utils.pad_shard_unpad, static_argnames=("b",))
      def add(params, a, *, b):
          return params * a + b
  
      x = np.arange(bs, dtype=dtype)
      y = add(5, x, b=10)
      chex.assert_type(y.dtype, x.dtype)
      np.testing.assert_allclose(np.float64(y), np.float64(5 * x + 10))
  
file_path: flax/tests/jax_utils_test.py
imports: from functools import partial
  import chex
  import jax
  import numpy
  import tensorflow
  from absl.testing import parameterized
  from flax import jax_utils
label_desc: Use image node_id 6 for calling the PadShardUnpadTest obj's underlying member method code with example usage: obj.test_static_argnames(dtype, bs) and returns: unknown