code: def test_jax_bfloat16():
      features = features_lib.FeaturesDict(
          {
              "data": features_lib.Tensor(
                  shape=(1,), dtype=tf.bfloat16, encoding="bytes"
              )
          }
      )
      data = {
          "data": jnp.array([6.0], dtype=jnp.bfloat16),
      }
      features.encode_example(data)
  
imports: from absl.testing import parameterized
  import jax.numpy
  import numpy
  import pytest
  import tensorflow
  from tensorflow_datasets import testing
  from tensorflow_datasets.core import features
label_desc: Use image node_id 2 for calling a global function with example usage: test_jax_bfloat16() without return types