code: def test_basics(self, dtype, bs):
      # Just tests that basic calling works without exploring caveats.
      @partial(jax_utils.pad_shard_unpad, static_argnums=())
      def add(a, b):
          return a + b
  
      x = np.arange(bs, dtype=dtype)
      y = add(x, 10 * x)
      chex.assert_type(y.dtype, x.dtype)
      np.testing.assert_allclose(np.float64(y), np.float64(x + 10 * x))
  
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 2 for calling the PadShardUnpadTest obj's underlying member method code with example usage: obj.test_basics(dtype, bs) and returns: unknown