code: def rotate_half(x):
      x1, x2 = jnp.split(x, 2, axis=-1)
      x = jnp.concatenate([-x2, x1], axis=-1)
      return x
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 5 for calling a global function with example usage: rotate_half(x) and returns: x