code: def test_gradients(self):
      ds_spec = {
          "text": types.TextSegment(),
          "target": types.CategoryLabel(),
      }
      out_spec = {
          "tokens": types.Tokens(),
          "embs": types.Embeddings(),
          "grads": types.Gradients(
              align="tokens",
              grad_for="embs",
              grad_target_field_key="target",
          ),
      }
      example = {"text": "hi", "target": "one"}
      output = {"tokens": ["hi"], "embs": [0.1, 0.2], "grads": [0.1]}
  
      grads = types.Gradients(
          align="tokens",
          grad_for="embs",
          grad_target_field_key="target",
      )
      embs = types.Embeddings()
      try:
          grads.validate_output(
              [0.1], out_spec, output, ds_spec, ds_spec, example
          )
          embs.validate_output(
              [0.1, 0.2], out_spec, output, ds_spec, ds_spec, example
          )
      except ValueError:
          self.fail("Raised unexpected error.")
  
      self.assertRaises(
          ValueError,
          grads.validate_output,
          ["bad"],
          out_spec,
          output,
          ds_spec,
          ds_spec,
          example,
      )
      self.assertRaises(
          ValueError,
          embs.validate_output,
          ["bad"],
          out_spec,
          output,
          ds_spec,
          ds_spec,
          example,
      )
  
      bad_grads = types.Gradients(
          align="text", grad_for="embs", grad_target_field_key="target"
      )
      self.assertRaises(
          ValueError,
          bad_grads.validate_output,
          [0.1],
          out_spec,
          output,
          ds_spec,
          ds_spec,
          example,
      )
      bad_grads = types.Gradients(
          align="tokens",
          grad_for="tokens",
          grad_target_field_key="target",
      )
      self.assertRaises(
          ValueError,
          bad_grads.validate_output,
          [0.1],
          out_spec,
          output,
          ds_spec,
          ds_spec,
          example,
      )
      bad_grads = types.Gradients(
          align="tokens", grad_for="embs", grad_target_field_key="bad"
      )
      self.assertRaises(
          ValueError,
          bad_grads.validate_output,
          [0.1],
          out_spec,
          output,
          ds_spec,
          ds_spec,
          example,
      )
  
imports: from collections.abc import Callable
  import os
  from typing import Any, Optional, Union
  from absl.testing import absltest
  from absl.testing import parameterized
  from etils import epath
  from lit_nlp.api import dtypes
  from lit_nlp.api import types
  import numpy
label_desc: Use image node_id 18 for calling the TypesTest obj's underlying member method code with example usage: obj.test_gradients() without return types